diff --git a/crates/nu-lsp/src/diagnostics.rs b/crates/nu-lsp/src/diagnostics.rs index faa0e6dd33..8e37191b0a 100644 --- a/crates/nu-lsp/src/diagnostics.rs +++ b/crates/nu-lsp/src/diagnostics.rs @@ -3,24 +3,30 @@ use lsp_types::{ notification::{Notification, PublishDiagnostics}, Diagnostic, DiagnosticSeverity, PublishDiagnosticsParams, Uri, }; -use miette::{IntoDiagnostic, Result}; +use miette::{miette, IntoDiagnostic, Result}; impl LanguageServer { pub(crate) fn publish_diagnostics_for_file(&mut self, uri: Uri) -> Result<()> { let mut engine_state = self.new_engine_state(); engine_state.generate_nu_constant(); - let Some((_, offset, working_set, file)) = self.parse_file(&mut engine_state, &uri, true) - else { + let Some((_, offset, working_set)) = self.parse_file(&mut engine_state, &uri, true) else { return Ok(()); }; let mut diagnostics = PublishDiagnosticsParams { - uri, + uri: uri.clone(), diagnostics: Vec::new(), version: None, }; + let docs = match self.docs.lock() { + Ok(it) => it, + Err(err) => return Err(miette!(err.to_string())), + }; + let file = docs + .get_document(&uri) + .ok_or_else(|| miette!("\nFailed to get document"))?; for err in working_set.parse_errors.iter() { let message = err.to_string(); diff --git a/crates/nu-lsp/src/goto.rs b/crates/nu-lsp/src/goto.rs index 78fb4c9067..72e576d674 100644 --- a/crates/nu-lsp/src/goto.rs +++ b/crates/nu-lsp/src/goto.rs @@ -33,7 +33,7 @@ impl LanguageServer { .text_document .uri .to_owned(); - let (working_set, id, _, _, _) = self + let (working_set, id, _, _) = self .parse_and_find( &mut engine_state, &path_uri, diff --git a/crates/nu-lsp/src/lib.rs b/crates/nu-lsp/src/lib.rs index 3e85723bee..b6759fe99e 100644 --- a/crates/nu-lsp/src/lib.rs +++ b/crates/nu-lsp/src/lib.rs @@ -1,5 +1,6 @@ #![doc = include_str!("../README.md")] use ast::find_id; +use crossbeam_channel::{Receiver, Sender}; use lsp_server::{Connection, IoThreads, Message, Response, ResponseError}; use lsp_textdocument::{FullTextDocument, TextDocuments}; use lsp_types::{ @@ -18,7 +19,7 @@ use nu_protocol::{ engine::{CachedFile, EngineState, Stack, StateWorkingSet}, DeclId, ModuleId, Span, Type, Value, VarId, }; -use std::collections::BTreeMap; +use std::{collections::BTreeMap, sync::Mutex}; use std::{ path::{Path, PathBuf}, str::FromStr, @@ -27,6 +28,7 @@ use std::{ }; use symbols::SymbolCache; use url::Url; +use workspace::{InternalMessage, RangePerDoc}; mod ast; mod diagnostics; @@ -47,13 +49,14 @@ pub enum Id { pub struct LanguageServer { connection: Connection, io_threads: Option, - docs: TextDocuments, + docs: Arc>, engine_state: EngineState, symbol_cache: SymbolCache, inlay_hints: BTreeMap>, workspace_folders: BTreeMap, // for workspace wide requests occurrences: BTreeMap>, + channels: Option<(Sender, Arc>)>, } pub fn path_to_uri(path: impl AsRef) -> Uri { @@ -92,12 +95,13 @@ impl LanguageServer { Ok(Self { connection, io_threads, - docs: TextDocuments::new(), + docs: Arc::new(Mutex::new(TextDocuments::new())), engine_state, symbol_cache: SymbolCache::new(), inlay_hints: BTreeMap::new(), workspace_folders: BTreeMap::new(), occurrences: BTreeMap::new(), + channels: None, }) } @@ -141,12 +145,19 @@ impl LanguageServer { self.initialize_workspace_folders(init_params)?; while !self.engine_state.signals().interrupted() { + // first check new messages from child thread + self.handle_internal_messages()?; + let msg = match self .connection .receiver .recv_timeout(Duration::from_secs(1)) { - Ok(msg) => msg, + Ok(msg) => { + // cancel execution if other messages received before job done + self.cancel_background_thread(); + msg + } Err(crossbeam_channel::RecvTimeoutError::Timeout) => { continue; } @@ -177,7 +188,9 @@ impl LanguageServer { Self::handle_lsp_request(request, |params| self.document_symbol(params)) } request::References::METHOD => { - Self::handle_lsp_request(request, |params| self.references(params)) + Self::handle_lsp_request(request, |params| { + self.references(params, 5000) + }) } request::WorkspaceSymbolRequest::METHOD => { Self::handle_lsp_request(request, |params| { @@ -224,6 +237,42 @@ impl LanguageServer { Ok(()) } + /// Send a cancel message to a running bg thread + pub fn cancel_background_thread(&mut self) { + if let Some((sender, _)) = &self.channels { + sender.send(true).ok(); + } + } + + /// Check results from background thread + pub fn handle_internal_messages(&mut self) -> Result { + let mut reset = false; + if let Some((_, receiver)) = &self.channels { + for im in receiver.try_iter() { + match im { + InternalMessage::RangeMessage(RangePerDoc { uri, ranges }) => { + self.occurrences.insert(uri, ranges); + } + InternalMessage::OnGoing(token, progress) => { + self.send_progress_report(token, progress, None)?; + } + InternalMessage::Finished(token) => { + reset = true; + self.send_progress_end(token, Some("Finished.".to_string()))?; + } + InternalMessage::Cancelled(token) => { + reset = true; + self.send_progress_end(token, Some("interrupted.".to_string()))?; + } + } + } + } + if reset { + self.channels = None; + } + Ok(reset) + } + pub fn new_engine_state(&self) -> EngineState { let mut engine_state = self.engine_state.clone(); let cwd = std::env::current_dir().expect("Could not get current working directory."); @@ -236,17 +285,24 @@ impl LanguageServer { engine_state: &'a mut EngineState, uri: &Uri, pos: Position, - ) -> Result<(StateWorkingSet<'a>, Id, Span, usize, &FullTextDocument)> { - let (block, file_offset, mut working_set, file) = self + ) -> Result<(StateWorkingSet<'a>, Id, Span, usize)> { + let (block, file_offset, mut working_set) = self .parse_file(engine_state, uri, false) .ok_or_else(|| miette!("\nFailed to parse current file"))?; + let docs = match self.docs.lock() { + Ok(it) => it, + Err(err) => return Err(miette!(err.to_string())), + }; + let file = docs + .get_document(uri) + .ok_or_else(|| miette!("\nFailed to get document"))?; let location = file.offset_at(pos) as usize + file_offset; let (id, span) = find_id(&block, &working_set, &location) .ok_or_else(|| miette!("\nFailed to find current name"))?; // add block to working_set for later references working_set.add_block(block); - Ok((working_set, id, span, file_offset, file)) + Ok((working_set, id, span, file_offset)) } pub fn parse_file<'a>( @@ -254,9 +310,10 @@ impl LanguageServer { engine_state: &'a mut EngineState, uri: &Uri, need_hints: bool, - ) -> Option<(Arc, usize, StateWorkingSet<'a>, &FullTextDocument)> { + ) -> Option<(Arc, usize, StateWorkingSet<'a>)> { let mut working_set = StateWorkingSet::new(engine_state); - let file = self.docs.get_document(uri)?; + let docs = self.docs.lock().ok()?; + let file = docs.get_document(uri)?; let file_path = uri_to_path(uri); let file_path_str = file_path.to_str()?; let contents = file.get_content(None).as_bytes(); @@ -270,7 +327,7 @@ impl LanguageServer { let file_inlay_hints = self.extract_inlay_hints(&working_set, &block, offset, file); self.inlay_hints.insert(uri.clone(), file_inlay_hints); } - Some((block, offset, working_set, file)) + Some((block, offset, working_set)) } fn get_location_by_span<'a>( @@ -285,10 +342,10 @@ impl LanguageServer { return None; } let target_uri = path_to_uri(path); - if let Some(doc) = self.docs.get_document(&target_uri) { + if let Some(file) = self.docs.lock().ok()?.get_document(&target_uri) { return Some(Location { uri: target_uri, - range: span_to_range(span, doc, cached_file.covered_span.start), + range: span_to_range(span, file, cached_file.covered_span.start), }); } else { // in case where the document is not opened yet, typically included by `nu -I` @@ -344,7 +401,7 @@ impl LanguageServer { .text_document .uri .to_owned(); - let (working_set, id, _, _, _) = self + let (working_set, id, _, _) = self .parse_and_find( &mut engine_state, &path_uri, @@ -525,7 +582,8 @@ impl LanguageServer { fn complete(&mut self, params: &CompletionParams) -> Option { let path_uri = params.text_document_position.text_document.uri.to_owned(); - let file = self.docs.get_document(&path_uri)?; + let docs = self.docs.lock().ok()?; + let file = docs.get_document(&path_uri)?; let mut completer = NuCompleter::new(Arc::new(self.engine_state.clone()), Arc::new(Stack::new())); diff --git a/crates/nu-lsp/src/notification.rs b/crates/nu-lsp/src/notification.rs index 83925230be..399d1237f6 100644 --- a/crates/nu-lsp/src/notification.rs +++ b/crates/nu-lsp/src/notification.rs @@ -17,8 +17,8 @@ impl LanguageServer { &mut self, notification: lsp_server::Notification, ) -> Option { - self.docs - .listen(notification.method.as_str(), ¬ification.params); + let mut docs = self.docs.lock().ok()?; + docs.listen(notification.method.as_str(), ¬ification.params); match notification.method.as_str() { DidOpenTextDocument::METHOD => { let params: DidOpenTextDocumentParams = @@ -57,7 +57,7 @@ impl LanguageServer { } } - fn send_progress_notification( + pub fn send_progress_notification( &self, token: ProgressToken, value: WorkDoneProgress, @@ -74,12 +74,13 @@ impl LanguageServer { .into_diagnostic() } - pub fn send_progress_begin(&self, token: ProgressToken, title: &str) -> Result<()> { + pub fn send_progress_begin(&self, token: ProgressToken, title: String) -> Result<()> { self.send_progress_notification( token, WorkDoneProgress::Begin(WorkDoneProgressBegin { - title: title.to_string(), + title, percentage: Some(0), + cancellable: Some(true), ..Default::default() }), ) @@ -95,8 +96,8 @@ impl LanguageServer { token, WorkDoneProgress::Report(WorkDoneProgressReport { message, + cancellable: Some(true), percentage: Some(percentage), - ..Default::default() }), ) } diff --git a/crates/nu-lsp/src/symbols.rs b/crates/nu-lsp/src/symbols.rs index 604c867e01..625a0250d3 100644 --- a/crates/nu-lsp/src/symbols.rs +++ b/crates/nu-lsp/src/symbols.rs @@ -272,7 +272,8 @@ impl LanguageServer { ) -> Option { let engine_state = self.new_engine_state(); let uri = params.text_document.uri.to_owned(); - self.symbol_cache.update(&uri, &engine_state, &self.docs); + let docs = self.docs.lock().ok()?; + self.symbol_cache.update(&uri, &engine_state, &docs); Some(DocumentSymbolResponse::Flat( self.symbol_cache.get_symbols_by_uri(&uri)?, )) @@ -284,7 +285,8 @@ impl LanguageServer { ) -> Option { if self.symbol_cache.any_dirty() { let engine_state = self.new_engine_state(); - self.symbol_cache.update_all(&engine_state, &self.docs); + let docs = self.docs.lock().ok()?; + self.symbol_cache.update_all(&engine_state, &docs); } Some(WorkspaceSymbolResponse::Flat( self.symbol_cache diff --git a/crates/nu-lsp/src/workspace.rs b/crates/nu-lsp/src/workspace.rs index f3ac318a2d..8161719e61 100644 --- a/crates/nu-lsp/src/workspace.rs +++ b/crates/nu-lsp/src/workspace.rs @@ -4,11 +4,13 @@ use std::{ collections::{BTreeMap, HashMap}, fs, path::{Path, PathBuf}, + sync::Arc, }; use crate::{ ast::find_reference_by_id, path_to_uri, span_to_range, uri_to_path, Id, LanguageServer, }; +use crossbeam_channel::{Receiver, Sender}; use lsp_server::{Message, Request, Response}; use lsp_types::{ Location, PrepareRenameResponse, ProgressToken, Range, ReferenceParams, RenameParams, @@ -16,9 +18,75 @@ use lsp_types::{ }; use miette::{miette, IntoDiagnostic, Result}; use nu_glob::{glob, Paths}; -use nu_protocol::{engine::StateWorkingSet, Span}; +use nu_protocol::{ + engine::{EngineState, StateWorkingSet}, + Span, +}; use serde_json::Value; +/// Message type indicating ranges of interest in each doc +#[derive(Debug)] +pub struct RangePerDoc { + pub uri: Uri, + pub ranges: Vec, +} + +/// Message sent from background thread to main +#[derive(Debug)] +pub enum InternalMessage { + RangeMessage(RangePerDoc), + Cancelled(ProgressToken), + Finished(ProgressToken), + OnGoing(ProgressToken, u32), +} + +fn find_nu_scripts_in_folder(folder_uri: &Uri) -> Result { + let path = uri_to_path(folder_uri); + if !path.is_dir() { + return Err(miette!("\nworkspace folder does not exist.")); + } + let pattern = format!("{}/**/*.nu", path.to_string_lossy()); + glob(&pattern).into_diagnostic() +} + +fn find_reference_in_file( + working_set: &mut StateWorkingSet, + file: &FullTextDocument, + fp: &Path, + id: &Id, +) -> Option> { + let fp_str = fp.to_str()?; + let block = parse( + working_set, + Some(fp_str), + file.get_content(None).as_bytes(), + false, + ); + let file_span = working_set.get_span_for_filename(fp_str)?; + let offset = file_span.start; + let mut references: Vec = find_reference_by_id(&block, working_set, id); + + // NOTE: for arguments whose declaration is in a signature + // which is not covered in the AST + if let Id::Variable(vid) = id { + let decl_span = working_set.get_variable(*vid).declaration_span; + if file_span.contains_span(decl_span) + && decl_span.end > decl_span.start + && !references.contains(&decl_span) + { + references.push(decl_span); + } + } + let occurs: Vec = references + .iter() + .map(|span| span_to_range(span, file, offset)) + .collect(); + + // add_block to avoid repeated parsing + working_set.add_block(block); + (!occurs.is_empty()).then_some(occurs) +} + impl LanguageServer { /// get initial workspace folders from initialization response pub fn initialize_workspace_folders(&mut self, init_params: Value) -> Result<()> { @@ -32,6 +100,8 @@ impl LanguageServer { Ok(()) } + /// The rename request only happens after the client received a `PrepareRenameResponse`, + /// and a new name typed in, could happen before ranges ready for all files in the workspace folder pub fn rename(&mut self, params: &RenameParams) -> Option { let new_name = params.new_name.to_owned(); // changes in WorkspaceEdit have mutable key @@ -59,31 +129,55 @@ impl LanguageServer { } /// Goto references response - /// TODO: WorkDoneProgress -> PartialResults - pub fn references(&mut self, params: &ReferenceParams) -> Option> { + /// # Arguments + /// - `timeout`: timeout in milliseconds, when timeout + /// 1. Respond with all ranges found so far + /// 2. Cancel the background thread + pub fn references(&mut self, params: &ReferenceParams, timeout: u128) -> Option> { self.occurrences = BTreeMap::new(); let mut engine_state = self.new_engine_state(); let path_uri = params.text_document_position.text_document.uri.to_owned(); - let (mut working_set, id, span, _, _) = self + let (working_set, id, span, _) = self .parse_and_find( &mut engine_state, &path_uri, params.text_document_position.position, ) .ok()?; - self.find_reference_in_workspace( - &mut working_set, - &path_uri, - id, - span, - params - .work_done_progress_params - .work_done_token - .to_owned() - .unwrap_or(ProgressToken::Number(1)), - "Finding references ...", - ) - .ok()?; + // have to clone it again in order to move to another thread + let mut engine_state = self.new_engine_state(); + engine_state.merge_delta(working_set.render()).ok()?; + let current_workspace_folder = self.get_workspace_folder_by_uri(&path_uri)?; + let token = params + .work_done_progress_params + .work_done_token + .to_owned() + .unwrap_or(ProgressToken::Number(1)); + self.channels = self + .find_reference_in_workspace( + engine_state, + current_workspace_folder, + id, + span, + token.clone(), + "Finding references ...".to_string(), + ) + .ok(); + // TODO: WorkDoneProgress -> PartialResults for quicker response + // currently not enabled by `lsp_types` but hackable in `server_capabilities` json + let time_start = std::time::Instant::now(); + loop { + if self.handle_internal_messages().ok()? { + break; + } + if time_start.elapsed().as_millis() > timeout { + self.send_progress_end(token, Some("Timeout".to_string())) + .ok()?; + self.cancel_background_thread(); + self.channels = None; + break; + } + } Some( self.occurrences .iter() @@ -108,7 +202,7 @@ impl LanguageServer { let mut engine_state = self.new_engine_state(); let path_uri = params.text_document.uri.to_owned(); - let (mut working_set, id, span, file_offset, file) = + let (working_set, id, span, file_offset) = self.parse_and_find(&mut engine_state, &path_uri, params.position)?; if let Id::Value(_) = id { @@ -119,6 +213,14 @@ impl LanguageServer { "\nDefinition not found.\nNot allowed to rename built-ins." )); } + + let docs = match self.docs.lock() { + Ok(it) => it, + Err(err) => return Err(miette!(err.to_string())), + }; + let file = docs + .get_document(&path_uri) + .ok_or_else(|| miette!("\nFailed to get document"))?; let range = span_to_range(&span, file, file_offset); let response = PrepareRenameResponse::Range(range); self.connection @@ -130,94 +232,112 @@ impl LanguageServer { })) .into_diagnostic()?; + // have to clone it again in order to move to another thread + let mut engine_state = self.new_engine_state(); + engine_state + .merge_delta(working_set.render()) + .into_diagnostic()?; + let current_workspace_folder = self + .get_workspace_folder_by_uri(&path_uri) + .ok_or_else(|| miette!("\nCurrent file is not in any workspace"))?; // now continue parsing on other files in the workspace - self.find_reference_in_workspace( - &mut working_set, - &path_uri, - id, - span, - ProgressToken::Number(0), - "Preparing rename ...", - ) + self.channels = self + .find_reference_in_workspace( + engine_state, + current_workspace_folder, + id, + span, + ProgressToken::Number(0), + "Preparing rename ...".to_string(), + ) + .ok(); + Ok(()) } fn find_reference_in_workspace( - &mut self, - working_set: &mut StateWorkingSet, - current_uri: &Uri, + &self, + engine_state: EngineState, + current_workspace_folder: WorkspaceFolder, id: Id, span: Span, token: ProgressToken, - message: &str, - ) -> Result<()> { - let current_workspace_folder = self - .get_workspace_folder_by_uri(current_uri) - .ok_or_else(|| miette!("\nCurrent file is not in any workspace"))?; - let scripts: Vec = Self::find_nu_scripts_in_folder(¤t_workspace_folder.uri)? - .filter_map(|p| p.ok()) - .collect(); - let len = scripts.len(); - + message: String, + ) -> Result<(Sender, Arc>)> { + let (data_sender, data_receiver) = crossbeam_channel::unbounded::(); + let (cancel_sender, cancel_receiver) = crossbeam_channel::bounded::(1); + let engine_state = Arc::new(engine_state); + let docs = self.docs.clone(); self.send_progress_begin(token.clone(), message)?; - for (i, fp) in scripts.iter().enumerate() { - let uri = path_to_uri(fp); - if let Some(file) = self.docs.get_document(&uri) { - Self::find_reference_in_file(working_set, file, fp, &id) - } else { - let bytes = fs::read(fp).into_diagnostic()?; - // skip if the file does not contain what we're looking for - let content_string = String::from_utf8(bytes).into_diagnostic()?; - let text_to_search = - String::from_utf8(working_set.get_span_contents(span).to_vec()) - .into_diagnostic()?; - if !content_string.contains(&text_to_search) { - continue; + + std::thread::spawn(move || -> Result<()> { + let mut working_set = StateWorkingSet::new(&engine_state); + let scripts: Vec = + match find_nu_scripts_in_folder(¤t_workspace_folder.uri) { + Ok(it) => it, + Err(_) => { + data_sender + .send(InternalMessage::Cancelled(token.clone())) + .ok(); + return Ok(()); + } } - let temp_file = FullTextDocument::new("nu".to_string(), 0, content_string); - Self::find_reference_in_file(working_set, &temp_file, fp, &id) + .filter_map(|p| p.ok()) + .collect(); + let len = scripts.len(); + + for (i, fp) in scripts.iter().enumerate() { + // std::thread::sleep(std::time::Duration::from_millis(500)); + // cancel the loop on cancellation message from main thread + if cancel_receiver.try_recv().is_ok() { + data_sender + .send(InternalMessage::Cancelled(token.clone())) + .into_diagnostic()?; + return Ok(()); + } + let percentage = (i * 100 / len) as u32; + let uri = path_to_uri(fp); + let docs = match docs.lock() { + Ok(it) => it, + Err(err) => return Err(miette!(err.to_string())), + }; + let file = if let Some(file) = docs.get_document(&uri) { + file + } else { + let bytes = match fs::read(fp) { + Ok(it) => it, + Err(_) => { + // continue on fs error + continue; + } + }; + // skip if the file does not contain what we're looking for + let content_string = String::from_utf8_lossy(&bytes); + let text_to_search = + String::from_utf8_lossy(working_set.get_span_contents(span)); + if !content_string.contains(text_to_search.as_ref()) { + // progress without any data + data_sender + .send(InternalMessage::OnGoing(token.clone(), percentage)) + .into_diagnostic()?; + continue; + } + &FullTextDocument::new("nu".to_string(), 0, content_string.into()) + }; + let _ = find_reference_in_file(&mut working_set, file, fp, &id).map(|ranges| { + data_sender + .send(InternalMessage::RangeMessage(RangePerDoc { uri, ranges })) + .ok(); + data_sender + .send(InternalMessage::OnGoing(token.clone(), percentage)) + .ok(); + }); } - .and_then(|range| self.occurrences.insert(uri, range)); - self.send_progress_report(token.clone(), (i * 100 / len) as u32, None)? - } - self.send_progress_end(token.clone(), Some("Done".to_string())) - } - - fn find_reference_in_file( - working_set: &mut StateWorkingSet, - file: &FullTextDocument, - fp: &Path, - id: &Id, - ) -> Option> { - let fp_str = fp.to_str()?; - let block = parse( - working_set, - Some(fp_str), - file.get_content(None).as_bytes(), - false, - ); - let file_span = working_set.get_span_for_filename(fp_str)?; - let offset = file_span.start; - let mut references: Vec = find_reference_by_id(&block, working_set, id); - - // NOTE: for arguments whose declaration is in a signature - // which is not covered in the AST - if let Id::Variable(vid) = id { - let decl_span = working_set.get_variable(*vid).declaration_span; - if file_span.contains_span(decl_span) - && decl_span.end > decl_span.start - && !references.contains(&decl_span) - { - references.push(decl_span); - } - } - let occurs: Vec = references - .iter() - .map(|span| span_to_range(span, file, offset)) - .collect(); - - // add_block to avoid repeated parsing - working_set.add_block(block); - (!occurs.is_empty()).then_some(occurs) + data_sender + .send(InternalMessage::Finished(token.clone())) + .into_diagnostic()?; + Ok(()) + }); + Ok((cancel_sender, Arc::new(data_receiver))) } fn get_workspace_folder_by_uri(&self, uri: &Uri) -> Option { @@ -231,15 +351,6 @@ impl LanguageServer { }) .cloned() } - - fn find_nu_scripts_in_folder(folder_uri: &Uri) -> Result { - let path = uri_to_path(folder_uri); - if !path.is_dir() { - return Err(miette!("\nworkspace folder does not exist.")); - } - let pattern = format!("{}/**/*.nu", path.to_string_lossy()); - glob(&pattern).into_diagnostic() - } } #[cfg(test)] @@ -525,7 +636,7 @@ mod tests { open_unchecked(&client_connection, script.clone()); - let message_num = 4; + let message_num = 5; let messages = send_rename_prepare_request(&client_connection, script.clone(), 3, 5, message_num); assert_eq!(messages.len(), message_num);