diff --git a/Cargo.toml b/Cargo.toml index 8becc8c..1357f13 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,7 +24,6 @@ tracing-appender = "0.2" tree-sitter-proto = "0.6" walkdir = "2.5" hard-xml = "1.41" -tempfile = "3.27" serde = { version = "1", features = ["derive"] } serde_json = "1.0" basic-toml = "0.1" @@ -34,6 +33,7 @@ const_format = "0.2" [dev-dependencies] insta = { version = "1.47", features = ["yaml", "redactions"] } +tempfile = "3.27.0" [profile.release] strip = "debuginfo" # Removes heavy debug data but keeps function names for panic logs @@ -42,3 +42,6 @@ codegen-units = 1 # Maximizes LLVM optimization passes [lints.clippy] pedantic = "warn" +large_futures = "warn" +large_types_passed_by_value = "warn" +inefficient_to_string = "warn" diff --git a/src/config/client.rs b/src/config/client.rs new file mode 100644 index 0000000..214bec5 --- /dev/null +++ b/src/config/client.rs @@ -0,0 +1,54 @@ +use async_lsp::lsp_types::{InitializeParams, WorkDoneProgressOptions}; + +/// A flat, optimized summary of client editor capabilities +/// compiled once during the LSP initialization handshake. +#[derive(Debug, Clone, Copy, Default)] +pub struct ClientCapabilitiesSummary { + /// True if the editor supports work done UI progress bars. + pub supports_work_done_progress: bool, + + /// True if the editor allows the server to dynamically register `didChangeWatchedFiles`. + pub supports_watched_files_dynamic_registration: bool, + + /// True if the editor can process modern directory-relative absolute glob patterns (LSP 3.17+). + pub supports_relative_patterns: bool, +} + +impl From<&InitializeParams> for ClientCapabilitiesSummary { + fn from(value: &InitializeParams) -> Self { + let supports_work_done_progress = Self::build_work_done_progress_options(value) + .work_done_progress + .unwrap_or_default(); + + let capabilies = &value.capabilities; + let watched_files_caps = capabilies + .workspace + .as_ref() + .and_then(|w| w.did_change_watched_files.as_ref()); + let supports_watched_files_dynamic_registration = watched_files_caps + .and_then(|cw| cw.dynamic_registration) + .unwrap_or_default(); + let supports_relative_patterns = watched_files_caps + .and_then(|cw| cw.relative_pattern_support) + .unwrap_or(false); + + Self { + supports_work_done_progress, + supports_watched_files_dynamic_registration, + supports_relative_patterns, + } + } +} + +impl ClientCapabilitiesSummary { + #[inline] + pub fn build_work_done_progress_options(params: &InitializeParams) -> WorkDoneProgressOptions { + let work_done_progress = params + .capabilities + .window + .as_ref() + .and_then(|w| w.work_done_progress); + + WorkDoneProgressOptions { work_done_progress } + } +} diff --git a/src/config/mod.rs b/src/config/mod.rs index 7190a33..07de25e 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -1,8 +1,12 @@ +use std::path::PathBuf; + +use serde::{Deserialize, Serialize}; + +pub(super) mod client; mod workspace; pub use workspace::WorkspaceProtoConfigs; - -use serde::{Deserialize, Serialize}; +pub use workspace::configure_workspace_folders; #[derive(Serialize, Deserialize, Debug, Clone, Default)] #[serde(default)] @@ -21,15 +25,15 @@ pub struct Config { #[derive(Serialize, Deserialize, Debug, Clone)] #[serde(default)] pub struct PathConfig { - pub clang_format: String, - pub protoc: String, + pub clang_format: PathBuf, + pub protoc: PathBuf, } impl Default for PathConfig { fn default() -> Self { Self { - clang_format: default_clang_format_path(), - protoc: default_protoc_path(), + clang_format: "clang-format".into(), + protoc: "protoc".into(), } } } @@ -44,11 +48,3 @@ pub struct RenameConfig { /// convention. pub chain_rpc_request_response: bool, } - -fn default_clang_format_path() -> String { - "clang-format".to_string() -} - -fn default_protoc_path() -> String { - "protoc".to_string() -} diff --git a/src/config/workspace.rs b/src/config/workspace.rs index ed0d12e..7575264 100644 --- a/src/config/workspace.rs +++ b/src/config/workspace.rs @@ -2,21 +2,21 @@ use std::{ collections::{HashMap, HashSet}, env, path::{Path, PathBuf}, + sync::Arc, }; use async_lsp::lsp_types::{Url, WorkspaceFolder}; +use futures::StreamExt; use pkg_config::Config; -use crate::formatter::ClangFormatter; - -use crate::config::ProtolsConfig; +use crate::{config::ProtolsConfig, formatter::clang::ClangFormatter, utils::OutermostPaths}; const CONFIG_FILE_NAMES: [&str; 2] = [".protols.toml", "protols.toml"]; pub struct WorkspaceProtoConfigs { workspaces: HashSet, configs: HashMap, - formatters: HashMap, + formatters: HashMap>, protoc_include_prefix: Vec, cli_include_paths: Vec, init_include_paths: Vec, @@ -47,34 +47,29 @@ impl WorkspaceProtoConfigs { } } - fn get_config_file_path(wpath: &PathBuf) -> Option { + pub(super) async fn get_config_file_path(wroot: &PathBuf) -> Option { for file in CONFIG_FILE_NAMES { - let p = Path::new(&wpath).join(file); - match std::fs::exists(&p) { - Ok(exists) if exists => return Some(p), - _ => {} + let p = Path::new(&wroot).join(file); + + if tokio::fs::try_exists(&p).await.unwrap_or_default() { + return Some(p); } } None } - pub fn add_workspace(&mut self, w: &WorkspaceFolder) { - let Ok(wpath) = w.uri.to_file_path() else { - return; - }; + pub fn add_workspace_folder(&mut self, data: ConfiguredWorkspaceFolder) { + self.workspaces.insert(data.uri.clone()); + self.configs.insert(data.uri.clone(), data.config); + self.formatters.insert(data.uri, Arc::new(data.formatter)); + } - let path = Self::get_config_file_path(&wpath).unwrap_or_default(); - let content = std::fs::read_to_string(path).unwrap_or_default(); + pub fn remove_workspace_folder(&mut self, wf: &WorkspaceFolder) { + let uri = &wf.uri; - let wr: ProtolsConfig = basic_toml::from_str(&content).unwrap_or_default(); - let fmt = ClangFormatter::new( - &wr.config.path.clang_format, - Some(wpath.to_str().expect("non-utf8 path")), - ); - - self.workspaces.insert(w.uri.clone()); - self.configs.insert(w.uri.clone(), wr); - self.formatters.insert(w.uri.clone(), fmt); + self.workspaces.remove(uri); + self.configs.remove(uri); + self.formatters.remove(uri); } pub fn get_config_for_uri(&self, u: &Url) -> Option<&ProtolsConfig> { @@ -82,9 +77,10 @@ impl WorkspaceProtoConfigs { .and_then(|w| self.configs.get(w)) } - pub fn get_formatter_for_uri(&self, u: &Url) -> Option<&ClangFormatter> { + pub fn get_formatter_for_uri(&self, u: &Url) -> Option> { self.get_workspace_for_uri(u) .and_then(|w| self.formatters.get(w)) + .cloned() } pub fn get_workspace_for_uri(&self, u: &Url) -> Option<&Url> { @@ -140,31 +136,97 @@ impl WorkspaceProtoConfigs { pub fn no_workspace_mode(&mut self) { let wr = ProtolsConfig::default(); - let rp = if cfg!(target_os = "windows") { - let mut d = String::from("C"); - if let Ok(cdir) = env::current_dir() - && let Some(drive) = cdir.components().next() - { - d = drive.as_os_str().to_string_lossy().to_string(); - } - format!("{d}://") - } else { - String::from("/") - }; - let uri = match Url::from_file_path(&rp) { - Err(err) => { - tracing::error!(?err, "failed to convert path: {rp} to Url"); + let root_path = env::current_dir() + .ok() + .and_then(|cdir| cdir.ancestors().last().map(Path::to_path_buf)) + .unwrap_or_else(|| { + if cfg!(target_os = "windows") { + r"C:\".into() + } else { + "/".into() + } + }); + let uri = match Url::from_file_path(&root_path) { + Err(error) => { + tracing::error!(?error, ?root_path, "failed to convert root path to Url"); return; } Ok(uri) => uri, }; - let fmt = ClangFormatter::new(&wr.config.path.clang_format, None); + let cmd = wr.config.path.clang_format.clone().into_os_string(); + let fmt = ClangFormatter::new(cmd, None); self.workspaces.insert(uri.clone()); self.configs.insert(uri.clone(), wr); - self.formatters.insert(uri.clone(), fmt); + self.formatters.insert(uri.clone(), Arc::new(fmt)); } + + pub fn get_all_indexing_paths(&self) -> OutermostPaths { + let workspace_paths = self + .workspaces + .iter() + .filter_map(|url| url.to_file_path().ok()); + + let include_paths = self + .workspaces + .iter() + .filter_map(|url| self.get_include_paths(url)) + .flatten(); + + OutermostPaths::from_iter(workspace_paths.chain(include_paths)) + } +} + +pub(crate) struct ConfiguredWorkspaceFolder { + pub uri: Url, + config: ProtolsConfig, + formatter: ClangFormatter, +} + +const MAX_CONCURRENT_DISK_SCANS: usize = 4; + +pub async fn configure_workspace_folders( + folders: &[WorkspaceFolder], +) -> Vec { + if folders.is_empty() { + return Vec::new(); + } + + futures::stream::iter(folders) + .map(configure_workspace_folder) + .buffer_unordered(MAX_CONCURRENT_DISK_SCANS) + .filter_map(|res| async move { res }) + .collect() + .await +} + +async fn configure_workspace_folder(wf: &WorkspaceFolder) -> Option { + let wfpath = wf.uri.to_file_path().ok()?; + let config = parse_config_file(&wfpath).await; + let cmd = config.config.path.clang_format.clone().into_os_string(); + + let formatter = ClangFormatter::new(cmd, Some(wfpath)); + + let uri = wf.uri.clone(); + Some(ConfiguredWorkspaceFolder { + uri, + config, + formatter, + }) +} + +#[inline] +async fn parse_config_file(path: &PathBuf) -> ProtolsConfig { + let Some(config_path) = WorkspaceProtoConfigs::get_config_file_path(path).await else { + return ProtolsConfig::default(); + }; + + let Ok(bytes) = tokio::fs::read(config_path).await else { + return ProtolsConfig::default(); + }; + + basic_toml::from_slice(&bytes).unwrap_or_default() } #[cfg(test)] @@ -174,24 +236,35 @@ mod test { use std::path::PathBuf; use tempfile::tempdir; - use super::{CONFIG_FILE_NAMES, WorkspaceProtoConfigs}; + use super::{CONFIG_FILE_NAMES, WorkspaceProtoConfigs, configure_workspace_folders}; - #[test] - fn test_get_for_workspace() { + #[tokio::test(flavor = "current_thread")] + async fn test_get_for_workspace() { let tmpdir = tempdir().expect("failed to create temp directory"); let tmpdir2 = tempdir().expect("failed to create temp2 directory"); let f = tmpdir.path().join("protols.toml"); std::fs::write(f, include_str!("input/protols-valid.toml")).unwrap(); + let folders = vec![ + WorkspaceFolder { + uri: Url::from_directory_path(tmpdir.path()).unwrap(), + name: "Test".to_owned(), + }, + WorkspaceFolder { + uri: Url::from_directory_path(tmpdir2.path()).unwrap(), + name: "Test2".to_owned(), + }, + ]; + + let configured_data = configure_workspace_folders(&folders).await; + + assert_eq!(configured_data.len(), folders.len()); + let mut ws = WorkspaceProtoConfigs::new(vec![], None); - ws.add_workspace(&WorkspaceFolder { - uri: Url::from_directory_path(tmpdir.path()).unwrap(), - name: "Test".to_string(), - }); - ws.add_workspace(&WorkspaceFolder { - uri: Url::from_directory_path(tmpdir2.path()).unwrap(), - name: "Test2".to_string(), - }); + + for data in configured_data { + ws.add_workspace_folder(data); + } let inworkspace = Url::from_file_path(tmpdir.path().join("foobar.proto")).unwrap(); let outworkspace = @@ -210,23 +283,33 @@ mod test { assert_yaml_snapshot!(ws.get_config_for_uri(&inworkspace2).unwrap()); } - #[test] - fn test_get_formatter_for_uri() { + #[tokio::test(flavor = "current_thread")] + async fn test_get_formatter_for_uri() { let tmpdir = tempdir().expect("failed to create temp directory"); let tmpdir2 = tempdir().expect("failed to create temp2 directory"); let f = tmpdir.path().join("protols.toml"); std::fs::write(f, include_str!("input/protols-valid.toml")).unwrap(); + let folders = vec![ + WorkspaceFolder { + uri: Url::from_directory_path(tmpdir.path()).unwrap(), + name: "Test".to_owned(), + }, + WorkspaceFolder { + uri: Url::from_directory_path(tmpdir2.path()).unwrap(), + name: "Test2".to_owned(), + }, + ]; + + let configured_data = configure_workspace_folders(&folders).await; + + assert_eq!(configured_data.len(), folders.len()); + let mut ws = WorkspaceProtoConfigs::new(vec![], None); - ws.add_workspace(&WorkspaceFolder { - uri: Url::from_directory_path(tmpdir.path()).unwrap(), - name: "Test".to_string(), - }); - ws.add_workspace(&WorkspaceFolder { - uri: Url::from_directory_path(tmpdir2.path()).unwrap(), - name: "Test2".to_string(), - }); + for data in configured_data { + ws.add_workspace_folder(data); + } let inworkspace = Url::from_file_path(tmpdir.path().join("foobar.proto")).unwrap(); let outworkspace = @@ -244,19 +327,27 @@ mod test { ); } - #[test] - fn test_loading_different_config_files() { + #[tokio::test(flavor = "current_thread")] + async fn test_loading_different_config_files() { let tmpdir = tempdir().expect("failed to create temp directory"); for file in CONFIG_FILE_NAMES { let f = tmpdir.path().join(file); std::fs::write(f, include_str!("input/protols-valid.toml")).unwrap(); - let mut ws = WorkspaceProtoConfigs::new(vec![], None); - ws.add_workspace(&WorkspaceFolder { + let folders = vec![WorkspaceFolder { uri: Url::from_directory_path(tmpdir.path()).unwrap(), - name: "Test".to_string(), - }); + name: "Test".to_owned(), + }]; + + let configured_data = configure_workspace_folders(&folders).await; + + assert_eq!(configured_data.len(), folders.len()); + + let mut ws = WorkspaceProtoConfigs::new(vec![], None); + for data in configured_data { + ws.add_workspace_folder(data); + } // check we really loaded the config file let workspace = Url::from_file_path(tmpdir.path().join("foobar.proto")).unwrap(); @@ -264,8 +355,8 @@ mod test { } } - #[test] - fn test_cli_include_paths() { + #[tokio::test(flavor = "current_thread")] + async fn test_cli_include_paths() { let tmpdir = tempdir().expect("failed to create temp directory"); let f = tmpdir.path().join("protols.toml"); std::fs::write(f, include_str!("input/protols-valid.toml")).unwrap(); @@ -275,11 +366,19 @@ mod test { // Set CLI include paths let cli_paths = vec![absolute_path.clone(), PathBuf::from("relative/path")]; - let mut ws = WorkspaceProtoConfigs::new(cli_paths, None); - ws.add_workspace(&WorkspaceFolder { + let folders = vec![WorkspaceFolder { uri: Url::from_directory_path(tmpdir.path()).unwrap(), - name: "Test".to_string(), - }); + name: "Test".to_owned(), + }]; + + let configured_data = configure_workspace_folders(&folders).await; + + assert_eq!(configured_data.len(), folders.len()); + + let mut ws = WorkspaceProtoConfigs::new(cli_paths, None); + for data in configured_data { + ws.add_workspace_folder(data); + } let inworkspace = Url::from_file_path(tmpdir.path().join("foobar.proto")).unwrap(); let include_paths = ws.get_include_paths(&inworkspace).unwrap(); @@ -292,8 +391,8 @@ mod test { assert!(include_paths.contains(&resolved_relative_path)); } - #[test] - fn test_init_include_paths() { + #[tokio::test(flavor = "current_thread")] + async fn test_init_include_paths() { let tmpdir = tempdir().expect("failed to create temp directory"); let f = tmpdir.path().join("protols.toml"); std::fs::write(f, include_str!("input/protols-valid.toml")).unwrap(); @@ -308,12 +407,20 @@ mod test { PathBuf::from("relative/init/path"), ]; + let folders = vec![WorkspaceFolder { + uri: Url::from_directory_path(tmpdir.path()).unwrap(), + name: "Test".to_owned(), + }]; + + let configured_data = configure_workspace_folders(&folders).await; + + assert_eq!(configured_data.len(), folders.len()); + let mut ws = WorkspaceProtoConfigs::new(cli_paths, None); ws.set_init_include_paths(init_paths); - ws.add_workspace(&WorkspaceFolder { - uri: Url::from_directory_path(tmpdir.path()).unwrap(), - name: "Test".to_string(), - }); + for data in configured_data { + ws.add_workspace_folder(data); + } let inworkspace = Url::from_file_path(tmpdir.path().join("foobar.proto")).unwrap(); let include_paths = ws.get_include_paths(&inworkspace).unwrap(); @@ -329,8 +436,8 @@ mod test { assert!(include_paths.contains(&cli_absolute_path)); } - #[test] - fn test_fallback_include_path() { + #[tokio::test(flavor = "current_thread")] + async fn test_fallback_include_path() { let tmpdir = tempdir().expect("failed to create temp directory"); let f = tmpdir.path().join("protols.toml"); std::fs::write(f, include_str!("input/protols-valid.toml")).unwrap(); @@ -342,12 +449,20 @@ mod test { PathBuf::from("relative/init/path"), ]; + let folders = vec![WorkspaceFolder { + uri: Url::from_directory_path(tmpdir.path()).unwrap(), + name: "Test".to_owned(), + }]; + + let configured_data = configure_workspace_folders(&folders).await; + + assert_eq!(configured_data.len(), folders.len()); + let mut ws = WorkspaceProtoConfigs::new(cli_paths, Some("fallback_path".into())); ws.set_init_include_paths(init_paths); - ws.add_workspace(&WorkspaceFolder { - uri: Url::from_directory_path(tmpdir.path()).unwrap(), - name: "Test".to_string(), - }); + for data in configured_data { + ws.add_workspace_folder(data); + } let inworkspace = Url::from_file_path(tmpdir.path().join("foobar.proto")).unwrap(); let include_paths = ws.get_include_paths(&inworkspace).unwrap(); diff --git a/src/document/diagnostics.rs b/src/document/diagnostics.rs index 818356a..e39f48a 100644 --- a/src/document/diagnostics.rs +++ b/src/document/diagnostics.rs @@ -10,7 +10,7 @@ impl ProtoDocument { range: r, severity: Some(DiagnosticSeverity::ERROR), source: Some(String::from("protols")), - message: "failed to find proto file".to_string(), + message: "failed to find proto file".to_owned(), ..Default::default() }) .collect() diff --git a/src/document/docsymbol.rs b/src/document/docsymbol.rs index 71a56a6..801029d 100644 --- a/src/document/docsymbol.rs +++ b/src/document/docsymbol.rs @@ -160,81 +160,85 @@ fn create_document_symbol(element: &ModelElement) -> Option { mod test { use async_lsp::lsp_types::{DocumentSymbol, Url}; use insta::assert_yaml_snapshot; + use tokio_util::sync::CancellationToken; - use crate::config::Config; - use crate::state::ProtoLanguageState; + use crate::server::state::{DocumentVersion, ProtoLanguageState}; - fn run_symbols_test(contents: &str, file_name: &str) -> Vec { + async fn run_symbols_test(contents: &str, file_name: &str) -> Vec { let uri = Url::parse(&format!("file:///virtual/{file_name}")).unwrap(); - let ipath = vec![]; - let mut state = ProtoLanguageState::new(); - state.upsert_file(&uri, contents, &ipath, 3, &Config::default(), false); + let state = ProtoLanguageState::new(); + state + .add_dirty_file_for_test(&uri, contents, DocumentVersion(1)) + .await; state - .get_document(&uri) + .get_document(&uri, CancellationToken::new()) + .await .map(|document| document.document_symbols()) - .unwrap_or_default() + .unwrap() } - #[test] - fn test_proto2_document_symbols() { + #[tokio::test(flavor = "current_thread")] + async fn test_proto2_document_symbols() { let contents = include_str!("input/syntax_variants/test_proto2.proto"); - let symbols = run_symbols_test(contents, "test_proto2.proto"); + let symbols = run_symbols_test(contents, "test_proto2.proto").await; assert_yaml_snapshot!("test_proto2_document_symbols", symbols); } - #[test] - fn test_proto3_document_symbols() { + #[tokio::test(flavor = "current_thread")] + async fn test_proto3_document_symbols() { let contents = include_str!("input/syntax_variants/test_proto3.proto"); - let symbols = run_symbols_test(contents, "test_proto3.proto"); + let symbols = run_symbols_test(contents, "test_proto3.proto").await; assert_yaml_snapshot!("test_proto3_document_symbols", symbols); } - #[test] - fn test_editions_document_symbols() { + #[tokio::test(flavor = "current_thread")] + async fn test_editions_document_symbols() { let contents = include_str!("input/syntax_variants/test_editions.proto"); - let symbols = run_symbols_test(contents, "test_editions.proto"); + let symbols = run_symbols_test(contents, "test_editions.proto").await; assert_yaml_snapshot!("test_editions_document_symbols", symbols); } - #[test] - fn test_package_duplicate_document_symbols() { + #[tokio::test(flavor = "current_thread")] + async fn test_package_duplicate_document_symbols() { let contents = include_str!("input/syntax_variants/test_package_duplicate.proto"); - let symbols = run_symbols_test(contents, "test_package_duplicate.proto"); + let symbols = run_symbols_test(contents, "test_package_duplicate.proto").await; assert_yaml_snapshot!("test_package_duplicate_document_symbols", symbols); } - #[test] - fn test_document_symbols_on_empty_and_minimal_file_safety() { + #[tokio::test(flavor = "current_thread")] + async fn test_document_symbols_on_empty_and_minimal_file_safety() { let uri = Url::parse("file:///virtual/empty_test.proto").unwrap(); - let ipath = vec![]; - let mut state = ProtoLanguageState::new(); - state.upsert_file(&uri, "", &ipath, 3, &Config::default(), false); + let state = ProtoLanguageState::new(); + state + .add_dirty_file_for_test(&uri, "", DocumentVersion(1)) + .await; let symbols = state - .get_document(&uri) + .get_document(&uri, CancellationToken::new()) + .await .map(|document| document.document_symbols()) .unwrap_or_default(); - assert!(symbols.is_empty()); + assert_eq!(symbols, [] as [DocumentSymbol; 0]); - let mut state_minimal = ProtoLanguageState::new(); - state_minimal.upsert_file( - &uri, - "syntax = \"proto3\";\npackage com.test;", - &ipath, - 3, - &Config::default(), - false, - ); + let state_minimal = ProtoLanguageState::new(); + state_minimal + .add_dirty_file_for_test( + &uri, + "syntax = \"proto3\";\npackage com.test;", + DocumentVersion(2), + ) + .await; let symbols_minimal = state_minimal - .get_document(&uri) + .get_document(&uri, CancellationToken::new()) + .await .map(|document| document.document_symbols()) .unwrap_or_default(); - assert!(symbols_minimal.is_empty()); + assert_eq!(symbols_minimal, [] as [DocumentSymbol; 0]); } } diff --git a/src/document/hover.rs b/src/document/hover.rs index b98a379..b8fc7d2 100644 --- a/src/document/hover.rs +++ b/src/document/hover.rs @@ -42,6 +42,7 @@ mod test { use async_lsp::lsp_types::{Hover, Position, Url}; use insta::assert_yaml_snapshot; use serde::Serialize; + use tokio_util::sync::CancellationToken; use crate::config::Config; use crate::model::ElementKind; @@ -53,7 +54,7 @@ mod test { hover: Hover, } - fn run_hover_test(contents: &str, file_name: &str) -> Vec { + async fn run_hover_test(contents: &str, file_name: &str) -> Vec { let uri = Url::parse(&format!("file:///virtual/{file_name}")).unwrap(); let ipath = vec![]; @@ -62,7 +63,11 @@ mod test { let mut hover_results = Vec::new(); - if let Some(parsed_document) = state.get_document(&uri) { + if let Some(parsed_document) = state + .get_document(&uri, CancellationToken::new()) + .await + .ok() + { let mut tested_positions = std::collections::HashSet::new(); for element in &parsed_document.elements { @@ -97,7 +102,8 @@ mod test { continue; } - if let Some(hover_result) = state.hover(&uri, pos) { + if let Ok(Some(hover)) = state.hover(&uri, pos, CancellationToken::new()).await + { let display_name = Some(element.meta.name.as_str()) .filter(|s| !s.is_empty()) .unwrap_or(match &element.kind { @@ -107,7 +113,7 @@ mod test { hover_results.push(HoverSnapshotEntry { target: format!("{display_name} [{context}]"), - hover: hover_result, + hover, }); } } @@ -117,36 +123,36 @@ mod test { hover_results } - #[test] - fn test_proto2_hover() { + #[tokio::test(flavor = "current_thread")] + async fn test_proto2_hover() { let contents = include_str!("input/syntax_variants/test_proto2.proto"); - let results = run_hover_test(contents, "test_proto2.proto"); + let results = run_hover_test(contents, "test_proto2.proto").await; assert_yaml_snapshot!("test_proto2_hover", results); } - #[test] - fn test_proto3_hover() { + #[tokio::test(flavor = "current_thread")] + async fn test_proto3_hover() { let contents = include_str!("input/syntax_variants/test_proto3.proto"); - let results = run_hover_test(contents, "test_proto3.proto"); + let results = run_hover_test(contents, "test_proto3.proto").await; assert_yaml_snapshot!("test_proto3_hover", results); } - #[test] - fn test_editions_hover() { + #[tokio::test(flavor = "current_thread")] + async fn test_editions_hover() { let contents = include_str!("input/syntax_variants/test_editions.proto"); - let results = run_hover_test(contents, "test_editions.proto"); + let results = run_hover_test(contents, "test_editions.proto").await; assert_yaml_snapshot!("test_editions_hover", results); } - #[test] - fn test_package_duplicate_hover() { + #[tokio::test(flavor = "current_thread")] + async fn test_package_duplicate_hover() { let contents = include_str!("input/syntax_variants/test_package_duplicate.proto"); - let results = run_hover_test(contents, "test_package_duplicate.proto"); + let results = run_hover_test(contents, "test_package_duplicate.proto").await; assert_yaml_snapshot!("test_package_duplicate_hover", results); } - #[test] - fn test_hover_on_empty_and_minimal_file_safety() { + #[tokio::test(flavor = "current_thread")] + async fn test_hover_on_empty_and_minimal_file_safety() { let uri = Url::parse("file:///virtual/empty_test.proto").unwrap(); let ipath = vec![]; @@ -157,7 +163,13 @@ mod test { line: 0, character: 0, }; - assert!(state.hover(&uri, pos).is_none()); + assert!( + state + .hover(&uri, pos, CancellationToken::new()) + .await + .ok() + .is_none() + ); let mut state_minimal = ProtoLanguageState::new(); state_minimal.upsert_file( @@ -173,6 +185,12 @@ mod test { line: 1, character: 5, }; - assert!(state_minimal.hover(&uri, pos_mid).is_none()); + assert!( + state_minimal + .hover(&uri, pos_mid, CancellationToken::new()) + .await + .ok() + .is_none() + ); } } diff --git a/src/document/mod.rs b/src/document/mod.rs index 50e5713..e9a77a1 100644 --- a/src/document/mod.rs +++ b/src/document/mod.rs @@ -2,8 +2,8 @@ mod parser; mod diagnostics; mod docsymbol; -mod hover; -mod rename; +// mod hover; +// mod rename; mod syntax; pub use parser::{ProtoDocument, ProtoParser}; diff --git a/src/document/syntax.rs b/src/document/syntax.rs index 432fca8..a3b0813 100644 --- a/src/document/syntax.rs +++ b/src/document/syntax.rs @@ -6,7 +6,7 @@ //! place we traverse the raw syntax tree directly. use async_lsp::lsp_types::{Diagnostic, DiagnosticSeverity}; -use tree_sitter::Node; +use tree_sitter::TreeCursor; use crate::utils::to_lsp_range; @@ -17,31 +17,60 @@ impl ProtoDocument { /// nodes. pub fn collect_parse_diagnostics(&self) -> Vec { let mut errors = Vec::new(); - collect_error_nodes(self.tree.root_node(), &mut errors); + let mut cursor = self.tree.walk(); + + loop { + let node = cursor.node(); + let is_missing = node.is_missing(); + + if node.is_error() || is_missing { + errors.push(Diagnostic { + range: to_lsp_range(node), + severity: Some(DiagnosticSeverity::ERROR), + source: Some("protols".into()), + message: if is_missing { + format!("Missing syntax element: {}", node.kind()) + } else { + "Syntax error".into() + }, + ..Default::default() + }); + + if advance_to_next(&mut cursor) { + continue; + } + + break; + } + + if cursor.goto_first_child() { + continue; + } + + if advance_to_next(&mut cursor) { + continue; + } + + break; + } errors - .into_iter() - .map(|n| Diagnostic { - range: to_lsp_range(n), - severity: Some(DiagnosticSeverity::ERROR), - source: Some("protols".to_string()), - message: "Syntax error".to_string(), - ..Default::default() - }) - .collect() } } -fn collect_error_nodes<'a>(n: Node<'a>, out: &mut Vec>) { - // Tree-sitter marks malformed regions with an `ERROR` node; these are the - // only raw-tree nodes we still inspect (they are absent from the metamodel). - if n.kind() == "ERROR" { - out.push(n); - } - let mut cursor = n.walk(); - for child in n.children(&mut cursor) { - collect_error_nodes(child, out); +/// Advances the cursor to the next node in the traversal, bypassing the current +/// node's children. +/// +/// Returns `true` if advanced to a sibling or a parent's sibling, or `false` if +/// the traversal is complete. +#[inline] +fn advance_to_next(cursor: &mut TreeCursor) -> bool { + while !cursor.goto_next_sibling() { + if !cursor.goto_parent() { + return false; + } } + true } #[cfg(test)] diff --git a/src/formatter/clang.rs b/src/formatter/clang.rs index f5e90dc..0009fb1 100644 --- a/src/formatter/clang.rs +++ b/src/formatter/clang.rs @@ -1,23 +1,17 @@ -#![allow(clippy::needless_late_init)] -use std::{ - borrow::Cow, - fs::File, - io::Write, - path::{Path, PathBuf}, - process::Command, -}; +use std::ffi::OsString; +use std::path::PathBuf; +use std::time::Duration; +use std::{borrow::Cow, process::Stdio}; use async_lsp::lsp_types::{Position, Range, TextEdit}; use hard_xml::XmlRead; use serde::Serialize; -use tempfile::{TempDir, tempdir}; - -use super::ProtoFormatter; +use tokio::io::AsyncWriteExt; +use tokio::process::Command; pub struct ClangFormatter { - pub path: String, - working_dir: Option, - temp_dir: TempDir, + pub path: OsString, + working_dir: Option, } #[derive(XmlRead, Serialize, PartialEq, Debug)] @@ -64,96 +58,126 @@ impl Replacement<'_> { }) } - fn as_text_edit(&self, content: &str) -> Option { + fn try_into_text_edit(self, content: &str) -> Option { Some(TextEdit { range: Range { start: Self::offset_to_position(self.offset, content)?, end: Self::offset_to_position(self.offset + self.length, content)?, }, - new_text: self.text.to_string(), + new_text: self.text.into_owned(), }) } } +const CLANG_FORMAT_TIMEOUT: Duration = Duration::from_secs(5); + impl ClangFormatter { - pub fn new(cmd: &str, wdir: Option<&str>) -> Self { + pub fn new(cmd: OsString, working_dir: Option) -> Self { Self { - temp_dir: tempdir().expect("faile to creat temp dir"), - path: cmd.to_owned(), - working_dir: wdir.map(ToOwned::to_owned), + path: cmd, + working_dir, } } - fn get_temp_file_path(&self, content: &str) -> Option { - let p = self.temp_dir.path().join("format-temp.proto"); - let mut file = File::create(p.clone()).ok()?; - file.write_all(content.as_ref()).ok()?; - Some(p) - } - - fn get_command(&self, f: &str, u: &Path) -> Option { - let mut c = Command::new(self.path.as_str()); - if let Some(wd) = &self.working_dir { - c.current_dir(wd.as_str()); - } - c.stdin(File::open(u).ok()?); - c.args([ - "--output-replacements-xml", - format!("--assume-filename={f}").as_str(), - ]); - Some(c) - } - - fn output_to_textedit(output: &str, content: &str) -> Option> { - let r = Replacements::from_str(output).ok()?; - let edits = r - .replacements - .into_iter() - .filter_map(|r| r.as_text_edit(content.as_ref())) - .collect(); - - Some(edits) - } -} - -impl ProtoFormatter for ClangFormatter { - fn format_document(&self, filename: &str, content: &str) -> Option> { - let p = self.get_temp_file_path(content)?; - let mut cmd = self.get_command(filename, p.as_ref())?; - let output = cmd.output().ok()?; - if !output.status.success() { - tracing::error!( - status = output.status.code(), - "failed to execute clang-format" - ); - return None; - } - Self::output_to_textedit(&String::from_utf8_lossy(&output.stdout), content) + /// # Cancellation safety + /// + /// This method is cancel safe. + pub async fn format_document(&self, filename: &str, content: &str) -> Option> { + let output = self.run_clang_format(filename, content, &[]).await?; + Self::output_to_textedit(&output, content) } - fn format_document_range( + /// # Cancellation safety + /// + /// This method is cancel safe. + pub async fn format_document_range( &self, r: &Range, filename: &str, content: &str, ) -> Option> { - let p = self.get_temp_file_path(content)?; let start = r.start.line + 1; let end = r.end.line + 1; + let line_arg = format!("{start}:{end}"); + let extra_args = ["--lines", &line_arg]; + let output = self - .get_command(filename, p.as_ref())? - .args(["--lines", format!("{start}:{end}").as_str()]) - .output() + .run_clang_format(filename, content, &extra_args) + .await?; + Self::output_to_textedit(&output, content) + } + + /// # Cancellation safety + /// + /// This method is cancel safe. + async fn run_clang_format( + &self, + filename: &str, + content: &str, + extra_args: &[&str], + ) -> Option { + let mut c = Command::new(&self.path); + + if let Some(wd) = &self.working_dir { + c.current_dir(wd); + } + + c.stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .args(["--output-replacements-xml", "--assume-filename", filename]) + .args(extra_args) + .kill_on_drop(true); + + let process = async { + let mut child = c.spawn()?; + + if let Some(mut stdin) = child.stdin.take() { + stdin.write_all(content.as_bytes()).await?; + } + + child.wait_with_output().await + }; + + let output_result = tokio::time::timeout(CLANG_FORMAT_TIMEOUT, process) + .await + .inspect_err(|_| { + tracing::error!( + filename, + timeout_ms = CLANG_FORMAT_TIMEOUT.as_millis(), + "clang-format execution timed out and was killed" + ); + }) .ok()?; - if !output.status.success() { + let output = output_result + .inspect_err(|error| tracing::error!(%error, "failed to run clang-format")) + .ok()?; + + if output.status.success() { + Some( + String::from_utf8(output.stdout) + .unwrap_or_else(|err| String::from_utf8_lossy(err.as_bytes()).into_owned()), + ) + } else { tracing::error!( status = output.status.code(), + error = %String::from_utf8_lossy(&output.stderr), "failed to execute clang-format" ); - return None; + None } - Self::output_to_textedit(&String::from_utf8_lossy(&output.stdout), content) + } + + fn output_to_textedit(output: &str, content: &str) -> Option> { + let r = Replacements::from_str(output).ok()?; + let edits = r + .replacements + .into_iter() + .filter_map(|r| r.try_into_text_edit(content)) + .collect(); + + Some(edits) } } @@ -227,14 +251,14 @@ mod test { " ); - let r = Replacements::from_str(&xml_output).unwrap(); + let mut r = Replacements::from_str(&xml_output).unwrap(); assert_eq!(r.replacements.len(), 1); - let replacement = &r.replacements[0]; - assert_eq!(replacement.offset, 134); + let replacement = r.replacements.remove(0); + assert_eq!(replacement.offset, offset); assert_eq!(replacement.length, 1); - let text_edit = replacement.as_text_edit(content).unwrap(); + let text_edit = replacement.try_into_text_edit(content).unwrap(); // The edit should be at line 1, character 77 (not 119) assert_eq!(text_edit.range.start.line, 1); assert_eq!(text_edit.range.start.character, 77); diff --git a/src/formatter/mod.rs b/src/formatter/mod.rs index cae7738..f537928 100644 --- a/src/formatter/mod.rs +++ b/src/formatter/mod.rs @@ -1,15 +1 @@ -mod clang; - -use async_lsp::lsp_types::{Range, TextEdit}; - -pub use clang::ClangFormatter; - -pub trait ProtoFormatter: Sized { - fn format_document(&self, filename: &str, content: &str) -> Option>; - fn format_document_range( - &self, - r: &Range, - filename: &str, - content: &str, - ) -> Option>; -} +pub mod clang; diff --git a/src/lsp.rs b/src/lsp.rs index 6264865..43d8a00 100644 --- a/src/lsp.rs +++ b/src/lsp.rs @@ -153,11 +153,12 @@ impl ProtoLanguageServer { ) -> BoxFuture<'static, Result<(), ResponseError>> { info!("Received shutdown request"); self.shutdown_received = true; + self.shutdown_cancel_token.cancel(); Box::pin(async move { Ok(()) }) } pub(super) fn hover( - &self, + &mut self, param: HoverParams, ) -> BoxFuture<'static, Result, ResponseError>> { let uri = param.text_document_position_params.text_document.uri; @@ -360,7 +361,7 @@ impl ProtoLanguageServer { } pub(super) fn document_symbol( - &self, + &mut self, params: DocumentSymbolParams, ) -> BoxFuture<'static, Result, ResponseError>> { let uri = params.text_document.uri; diff --git a/src/main.rs b/src/main.rs index c95ab23..74ad835 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,3 @@ -use std::time::Duration; - use async_lsp::LanguageClient; use async_lsp::client_monitor::ClientProcessMonitorLayer; use async_lsp::concurrency::ConcurrencyLayer; @@ -8,7 +6,8 @@ use async_lsp::server::LifecycleLayer; use async_lsp::tracing::TracingLayer; use clap::Parser; use cli::Cli; -use server::{ProtoLanguageServer, TickEvent}; +use server::ProtoLanguageServer; +use tokio_util::sync::CancellationToken; use tower::ServiceBuilder; use crate::transport::create_transport; @@ -19,11 +18,11 @@ mod docs; mod document; mod formatter; mod log; -mod lsp; +// mod lsp; mod model; mod protoc; mod server; -mod state; +// mod state; mod transport; mod utils; @@ -33,12 +32,24 @@ const FALLBACK_INCLUDE_PATH: Option<&str> = option_env!("FALLBACK_INCLUDE_PATH") async fn main() -> Result<(), transport::TransportError> { let cli = Cli::parse(); + let shutdown_token = CancellationToken::new(); + + tokio::spawn({ + let token = shutdown_token.clone(); + async move { + utils::wait_for_shutdown_signals().await; + token.cancel(); + } + }); + let (tx, mut rx) = tokio::sync::mpsc::channel(100); let (reload_handle, _log_guard) = log::install(tx); tracing::info!("server version: {}", env!("CARGO_PKG_VERSION")); tracing::info!("CLI include paths: {:?}", &cli.include_paths); + let shutdown_token_for_server = shutdown_token.child_token(); + let (server, _) = async_lsp::MainLoop::new_server(|client| { let mut log_client = client.clone(); @@ -59,33 +70,27 @@ async fn main() -> Result<(), transport::TransportError> { reload_handle, include_paths, fallback_include_path, + shutdown_token_for_server, ); - tokio::spawn({ - let client = client.clone(); - async move { - let mut interval = tokio::time::interval(Duration::from_secs(1)); - loop { - interval.tick().await; - if client.emit(TickEvent).is_err() { - break; - } - } - } - }); - ServiceBuilder::new() .layer(TracingLayer::default()) .layer(LifecycleLayer::default()) .layer(CatchUnwindLayer::default()) .layer(ConcurrencyLayer::default()) - .layer(ClientProcessMonitorLayer::new(client.clone())) + .layer(ClientProcessMonitorLayer::new(client)) .service(router) }); let (input, output) = create_transport(&cli).await?; - server.run_buffered(input, output).await?; + if shutdown_token + .run_until_cancelled(server.run_buffered(input, output)) + .await + .is_none() + { + tracing::info!("Graceful shutdown complete. Exiting..."); + } Ok(()) } diff --git a/src/model/extractor.rs b/src/model/extractor.rs index 793b6a7..1e56535 100644 --- a/src/model/extractor.rs +++ b/src/model/extractor.rs @@ -114,7 +114,7 @@ fn build_fqn( match parent_fqn { Some(fqn) => format!("{fqn}.{entity_name}"), - None if package_name.is_empty() => entity_name.to_string(), + None if package_name.is_empty() => entity_name.to_owned(), None => format!("{package_name}.{entity_name}"), } } @@ -228,7 +228,7 @@ impl ParsedMatch { Vec::new() }; - let name = kind.fqn().unwrap_or_default().to_string(); + let name = kind.fqn().unwrap_or_default().to_owned(); let full_fqn = build_fqn(context_stack, elements, &name, package_name); kind.set_fqn(full_fqn); diff --git a/src/model/extractor/handlers/import.rs b/src/model/extractor/handlers/import.rs index 9b6b3a3..3c21b86 100644 --- a/src/model/extractor/handlers/import.rs +++ b/src/model/extractor/handlers/import.rs @@ -31,7 +31,7 @@ pub(super) fn extract_import( }; path = node .utf8_text(source) - .map(|raw| raw.trim_matches(['"', '\'']).to_string()) + .map(|raw| raw.trim_matches(['"', '\'']).to_owned()) .inspect_err(trace_utf8_error) .ok(); } diff --git a/src/model/presentation.rs b/src/model/presentation.rs index 565e7a1..15aeb05 100644 --- a/src/model/presentation.rs +++ b/src/model/presentation.rs @@ -9,7 +9,7 @@ use std::fmt::Write; use async_lsp::lsp_types::{Position, SymbolKind}; -use crate::docs; +use crate::{docs, utils::is_position_inside_range}; use super::types::{CardinalityKind, ElementKind, ModelElement, TypeReference}; @@ -216,6 +216,41 @@ impl ModelElement { Some(()) } + + /// Evaluates if the cursor is positioned strictly over an embedded type + /// name token, returning its clean un-prefixed string identifier. + /// + /// This method maps geometric bounds of `TypeReference` entities (like + /// field types, map keys, or RPC signatures), ensuring that cardinality + /// tokens or stream modifiers do not affect type-level lookups. + pub fn inspect_nested_type_reference(&self, position: Position) -> Option<&str> { + match &self.kind { + ElementKind::Field { type_ref, .. } | ElementKind::OneofField { type_ref, .. } => { + is_position_inside_range(position, type_ref.range).then_some(type_ref.name.as_str()) + } + ElementKind::MapField { + key_type_ref, + value_type_ref, + .. + } => is_position_inside_range(position, key_type_ref.range) + .then_some(key_type_ref.name.as_str()) + .or_else(|| { + is_position_inside_range(position, value_type_ref.range) + .then_some(value_type_ref.name.as_str()) + }), + ElementKind::Rpc { + request_type_ref, + response_type_ref, + .. + } => is_position_inside_range(position, request_type_ref.range) + .then_some(request_type_ref.name.as_str()) + .or_else(|| { + is_position_inside_range(position, response_type_ref.range) + .then_some(response_type_ref.name.as_str()) + }), + _ => None, + } + } } impl ElementKind { diff --git a/src/protoc.rs b/src/protoc.rs index 247bb74..f505be1 100644 --- a/src/protoc.rs +++ b/src/protoc.rs @@ -1,44 +1,61 @@ -use std::process::Command; +use std::ffi::OsStr; +use std::path::Path; +use std::time::Duration; use async_lsp::lsp_types::{Diagnostic, DiagnosticSeverity, Range}; +use tokio::process::Command; use tree_sitter::Point; use crate::utils::to_lsp_position; -pub fn collect_diagnostics( - protoc_path: &str, - file_path: &str, - include_paths: &[String], -) -> Vec { +const PROTOC_TIMEOUT: Duration = Duration::from_secs(5); + +/// # Cancellation safety +/// +/// This method is cancel safe. +pub async fn collect_diagnostics( + protoc_path: &Path, + file_path: &Path, + include_paths: I, +) -> Vec +where + I: IntoIterator, + P: AsRef, +{ let mut cmd = Command::new(protoc_path); // Add include paths for path in include_paths { - cmd.arg("-I").arg(path); + cmd.arg("--proto_path").arg(path); } // Generate descriptor but discard its output - cmd.arg("-o") - .arg(if cfg!(windows) { "NUL" } else { "/dev/null" }); - - // Add the file to check - cmd.arg(file_path); + cmd.arg("--descriptor_set_out") + .arg(if cfg!(windows) { "NUL" } else { "/dev/null" }) + // Add the file to check + .arg(file_path) + .kill_on_drop(true); // Run protoc and capture output - match cmd.output() { - Ok(output) => { - if output.status.success() { - Vec::new() - } else { - let error = String::from_utf8_lossy(&output.stderr); - parse_protoc_output(&error) - } - } - Err(e) => { - tracing::error!(error=%e, "failed to run protoc"); - Vec::new() - } + let Ok(cmd_result) = tokio::time::timeout(PROTOC_TIMEOUT, cmd.output()).await else { + tracing::warn!( + timeout_ms = PROTOC_TIMEOUT.as_millis(), + "protoc execution timed out and was killed", + ); + return Vec::new(); + }; + + let Ok(output) = + cmd_result.inspect_err(|error| tracing::error!(%error, "failed to run protoc")) + else { + return Vec::new(); + }; + + if output.status.success() { + return Vec::new(); } + + parse_protoc_output(&String::from_utf8_lossy(&output.stderr)) } // Visible for testing @@ -47,30 +64,38 @@ fn parse_protoc_output(output: &str) -> Vec { for line in output.lines() { // Parse protoc error format: file:line:column: message - if let Some((file_info, message)) = line.split_once(": ") { - let parts: Vec<&str> = file_info.split(':').collect(); - if parts.len() >= 3 - && let (Ok(line), Ok(col)) = (parts[1].parse::(), parts[2].parse::()) - { - let point = Point { - row: (line - 1) as usize, - column: (col - 1) as usize, - }; - let diagnostic = Diagnostic { - range: Range { - start: to_lsp_position(point), - end: to_lsp_position(Point { - row: point.row, - column: point.column + 1, - }), - }, - severity: Some(DiagnosticSeverity::ERROR), - source: Some("protoc".to_string()), - message: message.to_string(), - ..Default::default() - }; - diagnostics.push(diagnostic); - } + let Some((file_info, message)) = line.split_once(": ") else { + continue; + }; + + let Some((rest, col_str)) = file_info.rsplit_once(':') else { + continue; + }; + + let Some((_file_name, line_str)) = rest.rsplit_once(':') else { + continue; + }; + + if let (Ok(line), Ok(col)) = (line_str.parse::(), col_str.parse::()) { + let point = Point { + row: (line - 1) as usize, + column: (col - 1) as usize, + }; + + let diagnostic = Diagnostic { + range: Range { + start: to_lsp_position(point), + end: to_lsp_position(Point { + row: point.row, + column: point.column + 1, + }), + }, + severity: Some(DiagnosticSeverity::ERROR), + source: Some("protoc".to_owned()), + message: message.to_owned(), + ..Default::default() + }; + diagnostics.push(diagnostic); } } @@ -87,7 +112,7 @@ mod test { let diags = parse_protoc_output(output); assert_eq!(diags.len(), 1); assert_eq!(diags[0].message, "Expected field name."); - assert_eq!(diags[0].source, Some("protoc".to_string())); + assert_eq!(diags[0].source, Some("protoc".to_owned())); assert_eq!(diags[0].severity, Some(DiagnosticSeverity::ERROR)); assert_eq!(diags[0].range.start.line, 4); assert_eq!(diags[0].range.start.character, 2); @@ -106,14 +131,14 @@ mod test { #[test] fn test_parse_protoc_output_empty() { - assert!(parse_protoc_output("").is_empty()); + assert_eq!(parse_protoc_output(""), [] as [Diagnostic; 0]); } #[test] fn test_parse_protoc_output_malformed_line() { let output = "not a valid protoc error line\n"; let diags = parse_protoc_output(output); - assert!(diags.is_empty()); + assert_eq!(diags, [] as [Diagnostic; 0]); } #[test] @@ -121,14 +146,14 @@ mod test { // Missing column number let output = "foo.proto:5: Expected field name.\n"; let diags = parse_protoc_output(output); - assert!(diags.is_empty()); + assert_eq!(diags, [] as [Diagnostic; 0]); } #[test] fn test_parse_protoc_output_non_numeric_line_col() { let output = "foo.proto:abc:def: some message\n"; let diags = parse_protoc_output(output); - assert!(diags.is_empty()); + assert_eq!(diags, [] as [Diagnostic; 0]); } #[test] @@ -141,4 +166,21 @@ mod test { "'Foo' is not defined. It could be a typo for 'Bar'." ); } + + #[test] + fn test_parse_protoc_output_windows_absolute_path() { + let row = 10u32; + let col = 5u32; + let output = + format!("C:\\Users\\Admin\\project\\foo.proto:{row}:{col}: Field type not found.\n"); + let diags = parse_protoc_output(&output); + + assert_eq!(diags.len(), 1); + assert_eq!(diags[0].message, "Field type not found."); + assert_eq!(diags[0].source, Some("protoc".to_owned())); + assert_eq!(diags[0].severity, Some(DiagnosticSeverity::ERROR)); + + assert_eq!(diags[0].range.start.line, row - 1); + assert_eq!(diags[0].range.start.character, col - 1); + } } diff --git a/src/server.rs b/src/server.rs index 880bbe2..d9aee11 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,9 +1,14 @@ +use std::ops::ControlFlow; +use std::path::PathBuf; +use std::sync::Arc; + use async_lsp::{ ClientSocket, lsp_types::{ notification::{ - DidChangeTextDocument, DidCreateFiles, DidDeleteFiles, DidOpenTextDocument, - DidRenameFiles, DidSaveTextDocument, Exit, SetTrace, + DidChangeTextDocument, DidChangeWatchedFiles, DidChangeWorkspaceFolders, + DidCloseTextDocument, DidCreateFiles, DidDeleteFiles, DidOpenTextDocument, + DidRenameFiles, DidSaveTextDocument, Exit, Initialized, SetTrace, }, request::{ Completion, DocumentSymbolRequest, Formatting, GotoDefinition, HoverRequest, @@ -13,40 +18,72 @@ use async_lsp::{ }, router::Router, }; -use std::{ops::ControlFlow, path::PathBuf}; +use tokio::sync::{RwLock, mpsc::Sender}; +use tokio_util::sync::{CancellationToken, DropGuard}; + +use crate::{config::WorkspaceProtoConfigs, log}; + +use notification::Notification; +use notification::worker::Worker; +use state::ProtoLanguageState; -use crate::{config::WorkspaceProtoConfigs, log, state::ProtoLanguageState}; +mod notification; +pub(crate) mod progress; +mod request; +pub(super) mod state; -pub struct TickEvent; pub struct ProtoLanguageServer { pub client: ClientSocket, - pub(crate) log_handle: log::LogReloadHandle, - pub counter: i32, pub state: ProtoLanguageState, - pub configs: WorkspaceProtoConfigs, - pub shutdown_received: bool, + pub configs: Arc>, + pub shutdown_token: CancellationToken, + notification_tx: Sender, } impl ProtoLanguageServer { + pub fn new( + client: ClientSocket, + log_handle: log::LogReloadHandle, + cli_include_paths: Vec, + fallback_include_path: Option, + shutdown_token: CancellationToken, + ) -> Self { + let configs = Arc::new(RwLock::new(WorkspaceProtoConfigs::new( + cli_include_paths, + fallback_include_path, + ))); + let state = ProtoLanguageState::new(); + let notification_tx = Worker::start( + state.clone(), + configs.clone(), + log_handle, + client.clone(), + shutdown_token.child_token(), + ); + + Self { + client, + state, + configs, + shutdown_token, + notification_tx, + } + } + pub fn new_router( client: ClientSocket, log_handle: log::LogReloadHandle, cli_include_paths: Vec, fallback_include_path: Option, + shutdown_token: CancellationToken, ) -> Router { - let mut router = Router::new(Self { + let mut router = Router::new(Self::new( client, log_handle, - counter: 0, - state: ProtoLanguageState::new(), - configs: WorkspaceProtoConfigs::new(cli_include_paths, fallback_include_path), - shutdown_received: false, - }); - - router.event::(|st, _| { - st.counter += 1; - ControlFlow::Continue(()) - }); + cli_include_paths, + fallback_include_path, + shutdown_token, + )); // Ignore any unknown notification. router.unhandled_notification(|_, notif| { @@ -55,29 +92,46 @@ impl ProtoLanguageServer { }); // Handling request - router.request::(ProtoLanguageServer::initialize); - router.request::(ProtoLanguageServer::shutdown); - router.request::(|st, params| st.hover(params)); - router.request::(ProtoLanguageServer::completion); - router.request::(ProtoLanguageServer::prepare_rename); - router.request::(ProtoLanguageServer::rename); - router.request::(ProtoLanguageServer::references); - router.request::(ProtoLanguageServer::definition); - router.request::(|st, params| st.document_symbol(params)); - router.request::(ProtoLanguageServer::workspace_symbol); - router.request::(ProtoLanguageServer::formatting); - router.request::(ProtoLanguageServer::range_formatting); + router + .request::(Self::initialize) + .request::(Self::shutdown) + // .request::(Self::hover) + // .request::(Self::completion) + // .request::(Self::prepare_rename) + // .request::(Self::rename) + // .request::(Self::references) + // .request::(Self::definition) + .request::(Self::document_symbol) + // .request::(Self::workspace_symbol) + .request::(Self::formatting) + .request::(Self::range_formatting); // Handling notification - router.notification::(ProtoLanguageServer::set_trace); - router.notification::(ProtoLanguageServer::did_save); - router.notification::(ProtoLanguageServer::did_open); - router.notification::(ProtoLanguageServer::did_change); - router.notification::(ProtoLanguageServer::did_create_files); - router.notification::(ProtoLanguageServer::did_rename_files); - router.notification::(ProtoLanguageServer::did_delete_files); - router.notification::(ProtoLanguageServer::exit); + router + .notification::(Self::handle_notification::) + .notification::(Self::handle_notification::) + .notification::(Self::handle_notification::) + .notification::(Self::handle_notification::) + .notification::( + Self::handle_notification::, + ) + .notification::(Self::handle_notification::) + .notification::(Self::handle_notification::) + .notification::(Self::handle_notification::) + .notification::(Self::handle_notification::) + .notification::( + Self::handle_notification::, + ) + .notification::( + Self::handle_notification::, + ) + .notification::(Self::handle_notification::); router } + + #[inline] + pub fn new_request_guard(&self) -> DropGuard { + self.shutdown_token.child_token().drop_guard() + } } diff --git a/src/server/notification.rs b/src/server/notification.rs new file mode 100644 index 0000000..da7b75c --- /dev/null +++ b/src/server/notification.rs @@ -0,0 +1,117 @@ +use std::fmt; +use std::ops::ControlFlow; + +use async_lsp::lsp_types::notification::{self, Notification as N}; +use async_lsp::{Result, lsp_types}; +use tokio::sync::mpsc::error::TrySendError; + +use super::ProtoLanguageServer; + +use worker::WORKER_MESSAGE_BUFFER; + +pub(super) mod worker; + +#[derive(Debug)] +pub(super) enum Notification { + Initialized { + params: lsp_types::InitializedParams, + }, + OpenDocument { + params: lsp_types::DidOpenTextDocumentParams, + }, + ChangeDocument { + params: lsp_types::DidChangeTextDocumentParams, + }, + SaveDocument { + params: lsp_types::DidSaveTextDocumentParams, + }, + CloseDocument { + params: lsp_types::DidCloseTextDocumentParams, + }, + CreateFiles { + params: lsp_types::CreateFilesParams, + }, + RenameFiles { + params: lsp_types::RenameFilesParams, + }, + DeleteFiles { + params: lsp_types::DeleteFilesParams, + }, + ChangeWorkspaceFolders { + params: lsp_types::DidChangeWorkspaceFoldersParams, + }, + ChangeWatchedFiles { + params: lsp_types::DidChangeWatchedFilesParams, + }, + SetTrace { + params: lsp_types::SetTraceParams, + }, + Exit, +} + +impl ProtoLanguageServer { + pub(super) fn handle_notification(&mut self, params: N::Params) -> ControlFlow> + where + N: lsp_types::notification::Notification, + N::Params: Into, + { + match self.notification_tx.try_send(params.into()) { + Err(TrySendError::Full(_)) => { + tracing::error!( + method = N::METHOD, + WORKER_MESSAGE_BUFFER, + "CRITICAL: Notification worker channel is FULL" + ); + + return ControlFlow::Break(Err(async_lsp::Error::Io(std::io::Error::new( + std::io::ErrorKind::WouldBlock, + format!("Notification channel overflowed on method: {}", N::METHOD), + )))); + } + Err(TrySendError::Closed(_)) => { + tracing::error!( + method = N::METHOD, + "CRITICAL: Notification worker actor crashed or dropped its receiver!" + ); + + return ControlFlow::Break(Err(async_lsp::Error::Io(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + "Notification worker crashed unexpectedly", + )))); + } + Ok(()) => {} + } + + if N::METHOD != lsp_types::notification::Exit::METHOD { + return ControlFlow::Continue(()); + } + + if self.shutdown_token.is_cancelled() { + return ControlFlow::Break(Ok(())); + } + + ControlFlow::Break(Err(async_lsp::Error::Protocol( + "Received 'exit' notification without a prior 'shutdown' request".to_owned(), + ))) + } +} + +impl fmt::Display for Notification { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let s = match self { + Self::ChangeDocument { .. } => notification::DidChangeTextDocument::METHOD, + Self::ChangeWatchedFiles { .. } => notification::DidChangeWatchedFiles::METHOD, + Self::ChangeWorkspaceFolders { .. } => notification::DidChangeWorkspaceFolders::METHOD, + Self::CloseDocument { .. } => notification::DidCloseTextDocument::METHOD, + Self::CreateFiles { .. } => notification::DidCreateFiles::METHOD, + Self::DeleteFiles { .. } => notification::DidDeleteFiles::METHOD, + Self::Exit => notification::Exit::METHOD, + Self::Initialized { .. } => notification::Initialized::METHOD, + Self::OpenDocument { .. } => notification::DidOpenTextDocument::METHOD, + Self::RenameFiles { .. } => notification::DidRenameFiles::METHOD, + Self::SaveDocument { .. } => notification::DidSaveTextDocument::METHOD, + Self::SetTrace { .. } => notification::SetTrace::METHOD, + }; + f.write_str(s) + } +} diff --git a/src/server/notification/worker.rs b/src/server/notification/worker.rs new file mode 100644 index 0000000..a5d0d3b --- /dev/null +++ b/src/server/notification/worker.rs @@ -0,0 +1,141 @@ +use std::sync::Arc; + +use async_lsp::ClientSocket; +use async_lsp::lsp_types::ProgressToken; +use tokio::sync::RwLock; +use tokio_util::sync::CancellationToken; + +use super::super::progress::{LspProgressChannel, OptionChannelExt}; +use super::Notification; +use crate::config::WorkspaceProtoConfigs; +use crate::log::LogReloadHandle; +use crate::server::state::ProtoLanguageState; +use crate::server::state::finalizer::FinalizerToken; +use crate::utils::OutermostPaths; + +mod change_document; +mod change_watched_files; +mod change_workspace_folders; +mod close_document; +mod create_files; +mod delete_files; +mod exit; +mod file_watchers; +mod initialized; +mod open_document; +mod rename_files; +mod save_document; +mod set_trace; + +struct Indexing { + cancel_token: CancellationToken, + finalized: Option, +} + +pub struct Worker { + state: ProtoLanguageState, + configs: Arc>, + client: ClientSocket, + log_handle: LogReloadHandle, + shutdown_token: CancellationToken, + indexing: Indexing, +} + +pub(super) const WORKER_MESSAGE_BUFFER: usize = 256; + +impl Worker { + pub fn start( + state: ProtoLanguageState, + configs: Arc>, + log_handle: LogReloadHandle, + client: ClientSocket, + shutdown_token: CancellationToken, + ) -> tokio::sync::mpsc::Sender { + let (tx, mut rx) = tokio::sync::mpsc::channel(WORKER_MESSAGE_BUFFER); + + let indexing = Indexing { + cancel_token: shutdown_token.child_token(), + finalized: None, + }; + let mut worker = Self { + state, + configs, + client, + log_handle, + shutdown_token, + indexing, + }; + + tokio::spawn(async move { + while let Some(event) = rx.recv().await { + if worker.shutdown_token.is_cancelled() && !matches!(event, Notification::Exit) { + tracing::info!( + %event, + "Skipping notification handling because shutdown is in progress." + ); + continue; + } + + match event { + Notification::Initialized { params } => worker.initialized(params).await, + Notification::OpenDocument { params } => worker.open_document(params).await, + Notification::ChangeDocument { params } => worker.change_document(params).await, + Notification::SaveDocument { params } => worker.save_document(params).await, + Notification::CloseDocument { params } => worker.close_document(params).await, + Notification::CreateFiles { params } => worker.create_files(params).await, + Notification::RenameFiles { params } => worker.rename_files(params).await, + Notification::DeleteFiles { params } => worker.delete_files(params).await, + Notification::ChangeWorkspaceFolders { params } => { + worker.change_workspace_folders(params).await; + } + Notification::ChangeWatchedFiles { params } => { + worker.change_watched_files(params).await; + } + Notification::SetTrace { params } => worker.set_trace(¶ms), + Notification::Exit => { + worker.exit(); + break; + } + } + } + }); + + tx + } + + async fn run_indexing_pipeline(&mut self, paths: OutermostPaths) { + self.indexing.cancel_token.cancel(); + if let Some(finalizer) = self.indexing.finalized.take() { + tracing::warn!( + "The indexing wasn't cancelled before running the new one. Cancelling..." + ); + let _ = finalizer.await; + } + + self.indexing.cancel_token = self.shutdown_token.child_token(); + self.state.set_indexing(true); + + let state = self.state.clone(); + let client = self.client.clone(); + let cancel_token_clone = self.indexing.cancel_token.clone(); + let progress_token = ProgressToken::String("protobuf-workspace-indexing".to_owned()); + + let task_handle = tokio::spawn(async move { + let capabilities = state.client_capabilities(); + let channel = + LspProgressChannel::try_create(capabilities, client, progress_token, false).await; + + channel.begin("Indexing workspace"); + + state + .index_folders(paths, cancel_token_clone, channel.reporter()) + .await; + + channel.end(); + + state.set_indexing(false); + }); + + self.indexing.finalized = Some(FinalizerToken::from_handle(task_handle)); + } +} diff --git a/src/server/notification/worker/change_document.rs b/src/server/notification/worker/change_document.rs new file mode 100644 index 0000000..0bb9d59 --- /dev/null +++ b/src/server/notification/worker/change_document.rs @@ -0,0 +1,69 @@ +use std::sync::Arc; + +use async_lsp::lsp_types::DidChangeTextDocumentParams; +use async_lsp::lsp_types::notification::{DidChangeTextDocument, Notification as N}; + +use super::{Notification, Worker}; +use crate::server::state::DocumentVersion; +use crate::server::state::cache::DocumentsGuardExt; +use crate::server::state::source::SourceEntry; + +impl Worker { + pub(super) async fn change_document(&mut self, params: DidChangeTextDocumentParams) { + let Some(change) = params.content_changes.into_iter().next() else { + tracing::warn!( + method = DidChangeTextDocument::METHOD, + uri = %params.text_document.uri, + version = params.text_document.version, + "Received change notification with empty content changes" + ); + + return; + }; + + let uri = params.text_document.uri; + let version = params.text_document.version; + let content = Arc::new(change.text); + + { + let (mut documents, mut sources) = self.state.write_lock_all().await; + + if !documents.set_dirty(&uri) { + tracing::warn!( + method = DidChangeTextDocument::METHOD, + %uri, + %version, + "Cache anomaly: received change notification for a document that is missing" + ); + } + + let source = sources.get_mut(&uri); + + if let Some(existing_source) = source { + existing_source.version = DocumentVersion(version); + existing_source.content = content; + } else { + tracing::error!( + method = DidChangeTextDocument::METHOD, + %uri, + %version, + "Cache anomaly: received change notification for a document that is missing from sources" + ); + + sources.insert( + uri, + SourceEntry { + version: DocumentVersion(version), + content, + }, + ); + } + } + } +} + +impl From for Notification { + fn from(params: DidChangeTextDocumentParams) -> Self { + Notification::ChangeDocument { params } + } +} diff --git a/src/server/notification/worker/change_watched_files.rs b/src/server/notification/worker/change_watched_files.rs new file mode 100644 index 0000000..31be061 --- /dev/null +++ b/src/server/notification/worker/change_watched_files.rs @@ -0,0 +1,71 @@ +use async_lsp::lsp_types::notification::{DidChangeWatchedFiles, Notification as N}; +use async_lsp::lsp_types::{DidChangeWatchedFilesParams, FileChangeType, FileEvent}; + +use super::{Notification, Worker}; +use crate::server::state::cache::DocumentsGuardExt; + +impl Worker { + pub(super) async fn change_watched_files(&mut self, params: DidChangeWatchedFilesParams) { + if params.changes.is_empty() { + tracing::info!( + method = DidChangeWatchedFiles::METHOD, + "Received the notification with empty changes list" + ); + + return; + } + + { + let (mut documents, mut sources) = self.state.write_lock_all().await; + + for FileEvent { typ, uri } in params.changes { + match typ { + FileChangeType::CREATED => { + if sources.remove(&uri).is_some() { + tracing::error!( + method = DidChangeWatchedFiles::METHOD, + %uri, + "CRITICAL Cache anomaly: received CREATED event for a URI that already has active content in sources" + ); + } + documents.set_dirty(&uri); + } + + FileChangeType::CHANGED => { + if sources.contains_key(&uri) { + continue; + } + + documents.set_dirty(&uri); + } + + FileChangeType::DELETED => { + documents.cancel_and_remove(&uri); + + if sources.remove(&uri).is_some() { + tracing::warn!( + method = DidChangeWatchedFiles::METHOD, + %uri, + "Cache anomaly: watched file was DELETED on disk while still active in sources" + ); + } + } + change_type => { + tracing::warn!( + method = DidChangeWatchedFiles::METHOD, + ?change_type, + %uri, + "Unexpecte file change type" + ); + } + } + } + } + } +} + +impl From for Notification { + fn from(params: DidChangeWatchedFilesParams) -> Self { + Notification::ChangeWatchedFiles { params } + } +} diff --git a/src/server/notification/worker/change_workspace_folders.rs b/src/server/notification/worker/change_workspace_folders.rs new file mode 100644 index 0000000..feb9c94 --- /dev/null +++ b/src/server/notification/worker/change_workspace_folders.rs @@ -0,0 +1,128 @@ +use std::path::PathBuf; + +use async_lsp::lsp_types::DidChangeWorkspaceFoldersParams; +use async_lsp::lsp_types::notification::{DidChangeWorkspaceFolders, Notification as N}; + +use super::{Notification, Worker}; +use crate::config::configure_workspace_folders; +use crate::server::state::ProtoLanguageState; +use crate::utils::OutermostPaths; + +impl Worker { + pub(super) async fn change_workspace_folders( + &mut self, + params: DidChangeWorkspaceFoldersParams, + ) { + if params.event.added.is_empty() && params.event.removed.is_empty() { + tracing::warn!( + method = DidChangeWorkspaceFolders::METHOD, + "Received change notification with empty content changes" + ); + + return; + } + + self.indexing.cancel_token.cancel(); + if let Some(finalized) = self.indexing.finalized.take() { + let _ = finalized.await; + } + + let configured_folders = configure_workspace_folders(¶ms.event.added).await; + + { + let mut configs = self.configs.write().await; + + for folder in ¶ms.event.removed { + configs.remove_workspace_folder(folder); + } + + for folder in configured_folders { + configs.add_workspace_folder(folder); + } + } + + if self.shutdown_token.is_cancelled() { + return; + } + + let new_outermost = self.configs.read().await.get_all_indexing_paths(); + let removed_paths: Vec<_> = params + .event + .removed + .iter() + .filter_map(|folder| folder.uri.to_file_path().ok()) + .collect(); + + if !removed_paths.is_empty() { + purge_removed_workspace_folders(&self.state, &removed_paths, &new_outermost).await; + } + + self.refresh_file_watchers(&new_outermost).await; + + self.run_indexing_pipeline(new_outermost).await; + } +} + +async fn purge_removed_workspace_folders( + state: &ProtoLanguageState, + removed_folders: &[PathBuf], + new_active_scope: &OutermostPaths, +) { + if removed_folders.is_empty() { + return; + } + + { + let (mut documents, mut sources) = state.write_lock_all().await; + + documents.retain(|url, entry| { + let Ok(file_path) = url.to_file_path() else { + return true; + }; + + let inside_removed_folder = removed_folders + .iter() + .any(|removed| file_path.starts_with(removed)); + + if !inside_removed_folder { + return true; + } + + let still_needed_by_include_path = new_active_scope + .iter() + .any(|active_outer| file_path.starts_with(active_outer)); + + if still_needed_by_include_path { + return true; + } + + entry.cancel(); + + false + }); + + sources.retain(|url, _| { + let Ok(file_path) = url.to_file_path() else { + return true; + }; + + let inside_removed_folder = removed_folders + .iter() + .any(|removed| file_path.starts_with(removed)); + + if !inside_removed_folder { + return true; + } + + new_active_scope + .iter() + .any(|active_outer| file_path.starts_with(active_outer)) + }); + } +} + +impl From for Notification { + fn from(params: DidChangeWorkspaceFoldersParams) -> Self { + Notification::ChangeWorkspaceFolders { params } + } +} diff --git a/src/server/notification/worker/close_document.rs b/src/server/notification/worker/close_document.rs new file mode 100644 index 0000000..9662619 --- /dev/null +++ b/src/server/notification/worker/close_document.rs @@ -0,0 +1,38 @@ +use async_lsp::lsp_types::DidCloseTextDocumentParams; +use async_lsp::lsp_types::notification::{DidCloseTextDocument, Notification as N}; + +use crate::server::state::cache::DocumentsGuardExt; + +use super::{Notification, Worker}; + +impl Worker { + pub(super) async fn close_document(&mut self, params: DidCloseTextDocumentParams) { + let uri = params.text_document.uri; + + { + let (mut documents, mut sources) = self.state.write_lock_all().await; + + if !documents.set_dirty(&uri) { + tracing::warn!( + method = DidCloseTextDocument::METHOD, + %uri, + "Cache anomaly: received close notification for a document that is missing" + ); + } + + if sources.remove(&uri).is_none() { + tracing::error!( + method = DidCloseTextDocument::METHOD, + %uri, + "Cache anomaly: received close notification for a document that was missing from sources" + ); + } + } + } +} + +impl From for Notification { + fn from(params: DidCloseTextDocumentParams) -> Self { + Notification::CloseDocument { params } + } +} diff --git a/src/server/notification/worker/create_files.rs b/src/server/notification/worker/create_files.rs new file mode 100644 index 0000000..d4abafc --- /dev/null +++ b/src/server/notification/worker/create_files.rs @@ -0,0 +1,49 @@ +use std::str::FromStr; + +use async_lsp::lsp_types::notification::{DidCreateFiles, Notification as N}; +use async_lsp::lsp_types::{CreateFilesParams, Url}; + +use crate::server::state::cache::DocumentsGuardExt; + +use super::{Notification, Worker}; + +impl Worker { + pub(super) async fn create_files(&mut self, params: CreateFilesParams) { + let uris: Vec<_> = params + .files + .iter() + .filter_map(|f| Url::from_str(f.uri.as_str()).ok()) + .collect(); + + if uris.is_empty() { + tracing::info!( + method = DidCreateFiles::METHOD, + "Received create notification with empty files list" + ); + + return; + } + + { + let (mut documents, mut sources) = self.state.write_lock_all().await; + + for uri in uris { + if sources.remove(&uri).is_some() { + tracing::error!( + method = DidCreateFiles::METHOD, + %uri, + "CRITICAL cache anomaly: received create notification for a URI that already has active content in sources!" + ); + } + + documents.set_dirty(&uri); + } + } + } +} + +impl From for Notification { + fn from(params: CreateFilesParams) -> Self { + Notification::CreateFiles { params } + } +} diff --git a/src/server/notification/worker/delete_files.rs b/src/server/notification/worker/delete_files.rs new file mode 100644 index 0000000..cbd1ddb --- /dev/null +++ b/src/server/notification/worker/delete_files.rs @@ -0,0 +1,49 @@ +use std::str::FromStr; + +use async_lsp::lsp_types::notification::{DidDeleteFiles, Notification as N}; +use async_lsp::lsp_types::{DeleteFilesParams, Url}; + +use crate::server::state::cache::DocumentsGuardExt; + +use super::{Notification, Worker}; + +impl Worker { + pub(super) async fn delete_files(&mut self, params: DeleteFilesParams) { + let uris: Vec<_> = params + .files + .iter() + .filter_map(|f| Url::from_str(f.uri.as_str()).ok()) + .collect(); + + if uris.is_empty() { + tracing::info!( + method = DidDeleteFiles::METHOD, + "Received delete notification with empty files list" + ); + + return; + } + + { + let (mut documents, mut sources) = self.state.write_lock_all().await; + + for uri in uris { + if sources.remove(&uri).is_some() { + tracing::warn!( + method = DidDeleteFiles::METHOD, + %uri, + "Cache anomaly: received delete notification for a document that was still active in sources" + ); + } + + documents.cancel_and_remove(&uri); + } + } + } +} + +impl From for Notification { + fn from(params: DeleteFilesParams) -> Self { + Notification::DeleteFiles { params } + } +} diff --git a/src/server/notification/worker/exit.rs b/src/server/notification/worker/exit.rs new file mode 100644 index 0000000..c29ef1d --- /dev/null +++ b/src/server/notification/worker/exit.rs @@ -0,0 +1,21 @@ +use super::{Notification, Worker}; + +impl Worker { + pub(super) fn exit(&mut self) { + if self.shutdown_token.is_cancelled() { + tracing::info!( + "Received exit notification after shutdown. Terminating background worker cleanly." + ); + } else { + tracing::warn!( + "Received exit notification WITHOUT shutdown! Process state is inconsistent." + ); + } + } +} + +impl From<()> for Notification { + fn from(_params: ()) -> Self { + Notification::Exit + } +} diff --git a/src/server/notification/worker/file_watchers.rs b/src/server/notification/worker/file_watchers.rs new file mode 100644 index 0000000..1fda1ba --- /dev/null +++ b/src/server/notification/worker/file_watchers.rs @@ -0,0 +1,101 @@ +use async_lsp::Error; +use async_lsp::lsp_types::notification::{DidChangeWatchedFiles, Notification as N}; +use async_lsp::lsp_types::request::{RegisterCapability, UnregisterCapability}; +use async_lsp::lsp_types::{ + DidChangeWatchedFilesRegistrationOptions, FileSystemWatcher, GlobPattern, OneOf, Registration, + RegistrationParams, RelativePattern, Unregistration, UnregistrationParams, Url, WatchKind, +}; + +use super::Worker; +use crate::config::client::ClientCapabilitiesSummary; +use crate::utils::OutermostPaths; + +pub(super) const FILE_WATCHER_ID: &str = "protols-proto-file-watcher"; + +impl Worker { + pub(super) async fn refresh_file_watchers(&self, paths: &OutermostPaths) { + let ClientCapabilitiesSummary { + supports_watched_files_dynamic_registration, + supports_relative_patterns, + .. + } = self.state.client_capabilities(); + + if !supports_watched_files_dynamic_registration { + return; + } + + let kind = Some(WatchKind::Create | WatchKind::Change | WatchKind::Delete); + + let watchers: Vec<_> = if supports_relative_patterns { + paths + .iter() + .filter_map(|p| Url::from_file_path(p).ok()) + .map(|u| FileSystemWatcher { + glob_pattern: GlobPattern::Relative(RelativePattern { + base_uri: OneOf::Right(u), + pattern: String::from("**/*.proto"), + }), + kind, + }) + .collect() + } else { + vec![FileSystemWatcher { + glob_pattern: GlobPattern::String(String::from("**/*.proto")), + kind, + }] + }; + + if watchers.is_empty() { + return; + } + + let unregister_params = UnregistrationParams { + unregisterations: vec![Unregistration { + id: String::from(FILE_WATCHER_ID), + method: DidChangeWatchedFiles::METHOD.to_owned(), + }], + }; + let trace_error = |error: &Error| { + tracing::warn!( + ?error, + "Client failed unregister dynamic file watcher capability." + ); + }; + + let _ = self + .client + .request::(unregister_params) + .await + .inspect_err(trace_error); + + let watcher_options = DidChangeWatchedFilesRegistrationOptions { watchers }; + let register_options = serde_json::to_value(watcher_options).unwrap_or_default(); + + let registration_params = RegistrationParams { + registrations: vec![Registration { + id: String::from(FILE_WATCHER_ID), + method: DidChangeWatchedFiles::METHOD.to_owned(), + register_options: Some(register_options), + }], + }; + + let trace_error = |error: &Error| { + tracing::error!( + ?error, + "CRITICAL. Client failed to register dynamic file watcher capability." + ); + }; + let trace_ok = |(): &()| { + tracing::info!( + "Dynamic file watchers successfully synchronized with current workspace." + ); + }; + + let _ = self + .client + .request::(registration_params) + .await + .inspect_err(trace_error) + .inspect(trace_ok); + } +} diff --git a/src/server/notification/worker/initialized.rs b/src/server/notification/worker/initialized.rs new file mode 100644 index 0000000..5512ff9 --- /dev/null +++ b/src/server/notification/worker/initialized.rs @@ -0,0 +1,24 @@ +use async_lsp::lsp_types::InitializedParams; + +use super::{Notification, Worker}; + +impl Worker { + pub(super) async fn initialized(&mut self, _params: InitializedParams) { + let paths = self.configs.read().await.get_all_indexing_paths(); + + if paths.is_empty() { + tracing::info!("No indexing paths found. Skipping initial workspace indexing."); + return; + } + + self.refresh_file_watchers(&paths).await; + + self.run_indexing_pipeline(paths).await; + } +} + +impl From for Notification { + fn from(params: InitializedParams) -> Self { + Notification::Initialized { params } + } +} diff --git a/src/server/notification/worker/open_document.rs b/src/server/notification/worker/open_document.rs new file mode 100644 index 0000000..fba79b5 --- /dev/null +++ b/src/server/notification/worker/open_document.rs @@ -0,0 +1,45 @@ +use std::sync::Arc; + +use async_lsp::lsp_types::DidOpenTextDocumentParams; + +use crate::server::state::DocumentVersion; +use crate::server::state::cache::DocumentsGuardExt; +use crate::server::state::source::SourceEntry; +use crate::server::state::unique_uris::UniqueUris; + +use super::{Notification, Worker}; + +impl Worker { + pub(super) async fn open_document(&mut self, params: DidOpenTextDocumentParams) { + let uri = params.text_document.uri; + let version = params.text_document.version; + let content = Arc::new(params.text_document.text); + + { + let (mut documents, mut sources) = self.state.write_lock_all().await; + + documents.set_dirty(&uri); + + sources.insert( + uri.clone(), + SourceEntry { + version: DocumentVersion(version), + content, + }, + ); + } + + let cancel_token = self.shutdown_token.child_token(); + + let _ = self + .state + .query_documents_batch(UniqueUris::new(uri), cancel_token, None, false) + .await; + } +} + +impl From for Notification { + fn from(params: DidOpenTextDocumentParams) -> Self { + Notification::OpenDocument { params } + } +} diff --git a/src/server/notification/worker/rename_files.rs b/src/server/notification/worker/rename_files.rs new file mode 100644 index 0000000..ef708ac --- /dev/null +++ b/src/server/notification/worker/rename_files.rs @@ -0,0 +1,58 @@ +use std::str::FromStr; + +use async_lsp::lsp_types::{RenameFilesParams, Url}; + +use super::{Notification, Worker}; +use crate::server::state::cache::DocumentsGuardExt; +use crate::server::state::source::SourceEntry; + +struct RenameTask { + old_uri: Url, + new_uri: Url, + extracted_source: Option, +} + +impl Worker { + pub(super) async fn rename_files(&mut self, params: RenameFilesParams) { + let mut tasks: Vec = params + .files + .iter() + .filter_map(|r| { + let old_uri = Url::from_str(&r.old_uri).ok()?; + let new_uri = Url::from_str(&r.new_uri).ok()?; + Some(RenameTask { + old_uri, + new_uri, + extracted_source: None, + }) + }) + .collect(); + + if tasks.is_empty() { + return; + } + + { + let (mut documents, mut sources) = self.state.write_lock_all().await; + + for task in &mut tasks { + documents.cancel_and_remove(&task.old_uri); + task.extracted_source = sources.remove(&task.old_uri); + } + + for task in tasks { + documents.set_dirty(&task.new_uri); + + if let Some(src) = task.extracted_source { + sources.insert(task.new_uri, src); + } + } + } + } +} + +impl From for Notification { + fn from(params: RenameFilesParams) -> Self { + Notification::RenameFiles { params } + } +} diff --git a/src/server/notification/worker/save_document.rs b/src/server/notification/worker/save_document.rs new file mode 100644 index 0000000..3166ef2 --- /dev/null +++ b/src/server/notification/worker/save_document.rs @@ -0,0 +1,15 @@ +use async_lsp::lsp_types::DidSaveTextDocumentParams; + +use super::{Notification, Worker}; + +impl Worker { + pub(super) async fn save_document(&mut self, params: DidSaveTextDocumentParams) { + todo!("Add protoc diagnostic update"); + } +} + +impl From for Notification { + fn from(params: DidSaveTextDocumentParams) -> Self { + Notification::SaveDocument { params } + } +} diff --git a/src/server/notification/worker/set_trace.rs b/src/server/notification/worker/set_trace.rs new file mode 100644 index 0000000..87cf233 --- /dev/null +++ b/src/server/notification/worker/set_trace.rs @@ -0,0 +1,16 @@ +use async_lsp::lsp_types::SetTraceParams; + +use super::{Notification, Worker}; +use crate::log; + +impl Worker { + pub(super) fn set_trace(&mut self, params: &SetTraceParams) { + log::update_level(&self.log_handle, params.value); + } +} + +impl From for Notification { + fn from(params: SetTraceParams) -> Self { + Notification::SetTrace { params } + } +} diff --git a/src/server/progress.rs b/src/server/progress.rs new file mode 100644 index 0000000..ec65131 --- /dev/null +++ b/src/server/progress.rs @@ -0,0 +1,156 @@ +use async_lsp::ClientSocket; +use async_lsp::lsp_types::WorkDoneProgressCreateParams; +use async_lsp::lsp_types::request::{Request, WorkDoneProgressCreate}; +use async_lsp::lsp_types::{ + ProgressParams, ProgressParamsValue, ProgressToken, WorkDoneProgress, WorkDoneProgressBegin, + WorkDoneProgressEnd, WorkDoneProgressReport, notification::Progress, +}; + +use crate::config::client::ClientCapabilitiesSummary; + +pub struct LspProgressChannel { + client: ClientSocket, + token: ProgressToken, + cancellable: bool, +} + +impl LspProgressChannel { + pub async fn try_create( + client_capabilities: ClientCapabilitiesSummary, + client: ClientSocket, + token: ProgressToken, + cancellable: bool, + ) -> Option { + if !client_capabilities.supports_work_done_progress { + return None; + } + + let request_params = WorkDoneProgressCreateParams { + token: token.clone(), + }; + + client + .request::(request_params) + .await + .inspect_err(|error| { + tracing::warn!( + ?error, + method = WorkDoneProgressCreate::METHOD, + "Client advertised the method support, but failed to allocate token: {:?}", + token + ); + }) + .ok()?; + + Some(Self { + client, + token, + cancellable, + }) + } + + pub fn reporter(&self) -> LspProgressReporter { + LspProgressReporter { + client: self.client.clone(), + token: self.token.clone(), + min_pct: 0, + max_pct: 100, + cancellable: self.cancellable, + } + } + + pub fn begin(&self, title: &str) { + let _ = self.client.notify::(ProgressParams { + token: self.token.clone(), + value: ProgressParamsValue::WorkDone(WorkDoneProgress::Begin(WorkDoneProgressBegin { + title: title.to_string(), + cancellable: Some(self.cancellable), + message: None, + percentage: Some(0), + })), + }); + } +} + +impl Drop for LspProgressChannel { + fn drop(&mut self) { + let _ = self.client.notify::(ProgressParams { + token: self.token.clone(), + value: ProgressParamsValue::WorkDone(WorkDoneProgress::End(WorkDoneProgressEnd { + message: None, + })), + }); + } +} + +pub trait OptionChannelExt { + fn begin(&self, title: &str); + fn reporter(&self) -> Option; + fn end(self); +} + +impl OptionChannelExt for Option { + fn begin(&self, title: &str) { + if let Some(channel) = self { + channel.begin(title); + } + } + + fn reporter(&self) -> Option { + self.as_ref().map(LspProgressChannel::reporter) + } + + fn end(self) {} +} + +#[derive(Clone)] +pub struct LspProgressReporter { + client: ClientSocket, + token: ProgressToken, + cancellable: bool, + min_pct: u32, + max_pct: u32, +} + +impl LspProgressReporter { + pub fn with_range(&self, min_pct: u32, max_pct: u32) -> Self { + Self { + client: self.client.clone(), + token: self.token.clone(), + cancellable: self.cancellable, + min_pct, + max_pct, + } + } + + pub fn report(&self, percentage: u32, message: &str) { + let scaled_pct = self.min_pct + (self.max_pct - self.min_pct) * percentage / 100; + let _ = self.client.notify::(ProgressParams { + token: self.token.clone(), + value: ProgressParamsValue::WorkDone(WorkDoneProgress::Report( + WorkDoneProgressReport { + cancellable: Some(self.cancellable), + message: Some(message.to_string()), + percentage: Some(scaled_pct), + }, + )), + }); + } +} + +pub trait OptionReporterExt { + fn report(&self, percentage: u32, message: &str); + fn with_range(&self, min_pct: u32, max_pct: u32) -> Self; +} + +impl OptionReporterExt for Option { + fn report(&self, percentage: u32, message: &str) { + if let Some(reporter) = self { + reporter.report(percentage, message); + } + } + + fn with_range(&self, min_pct: u32, max_pct: u32) -> Self { + self.as_ref().map(|r| r.with_range(min_pct, max_pct)) + } +} diff --git a/src/server/request.rs b/src/server/request.rs new file mode 100644 index 0000000..82476f6 --- /dev/null +++ b/src/server/request.rs @@ -0,0 +1,4 @@ +mod document_symbol; +mod formatting; +mod initialize; +mod shutdown; diff --git a/src/server/request/document_symbol.rs b/src/server/request/document_symbol.rs new file mode 100644 index 0000000..8c2f6b5 --- /dev/null +++ b/src/server/request/document_symbol.rs @@ -0,0 +1,26 @@ +use async_lsp::ResponseError; +use async_lsp::lsp_types::{DocumentSymbolParams, DocumentSymbolResponse}; +use futures::FutureExt; +use futures::future::BoxFuture; + +use crate::ProtoLanguageServer; + +impl ProtoLanguageServer { + pub(in crate::server) fn document_symbol( + &mut self, + params: DocumentSymbolParams, + ) -> BoxFuture<'static, Result, ResponseError>> { + let state = self.state.clone(); + let uri = params.text_document.uri; + let guard = self.new_request_guard(); + + async move { + let document = state.get_document(&uri, guard.token().clone()).await?; + let symbols = document.document_symbols(); + let response = DocumentSymbolResponse::Nested(symbols); + + Ok(Some(response)) + } + .boxed() + } +} diff --git a/src/server/request/formatting.rs b/src/server/request/formatting.rs new file mode 100644 index 0000000..b3dbc83 --- /dev/null +++ b/src/server/request/formatting.rs @@ -0,0 +1,87 @@ +use std::sync::Arc; + +use async_lsp::lsp_types::{ + DocumentFormattingParams, DocumentRangeFormattingParams, TextEdit, Url, +}; +use async_lsp::{ErrorCode, ResponseError}; +use futures::FutureExt; +use futures::future::BoxFuture; + +use crate::ProtoLanguageServer; +use crate::formatter::clang::ClangFormatter; + +impl ProtoLanguageServer { + pub(in crate::server) fn formatting( + &mut self, + params: DocumentFormattingParams, + ) -> BoxFuture<'static, Result>, ResponseError>> { + let uri = params.text_document.uri; + + self.execute_formatting_pipeline(uri.clone(), move |formatter, text| async move { + formatter.format_document(uri.path(), text.as_str()).await + }) + } + + pub(in crate::server) fn range_formatting( + &mut self, + params: DocumentRangeFormattingParams, + ) -> BoxFuture<'static, Result>, ResponseError>> { + let range = params.range; + let uri = params.text_document.uri; + + self.execute_formatting_pipeline(uri.clone(), move |formatter, text| async move { + formatter + .format_document_range(&range, uri.path(), text.as_str()) + .await + }) + } + + fn execute_formatting_pipeline( + &self, + uri: Url, + f: F, + ) -> BoxFuture<'static, Result>, ResponseError>> + where + F: FnOnce(Arc, Arc) -> Fut + Send + 'static, + Fut: futures::Future>> + Send + 'static, + { + let state = self.state.clone(); + let config = self.configs.clone(); + let guard = self.new_request_guard(); + + let uri_clone = uri.clone(); + let formatter = async move { + config + .read() + .await + .get_formatter_for_uri(&uri_clone) + .ok_or_else(|| ResponseError::new(ErrorCode::INTERNAL_ERROR, "Formatter not found")) + }; + + let text = async move { + state.get_document_text(&uri).await.ok_or_else(|| { + ResponseError::new(ErrorCode::INTERNAL_ERROR, "Document text not found") + }) + }; + + let formatting_pipeline = + futures::future::try_join(formatter, text).then(move |res| async move { + let (formatter, text) = res?; + Ok(f(formatter, text).await) + }); + + async move { + guard + .token() + .run_until_cancelled(formatting_pipeline) + .await + .unwrap_or_else(|| { + Err(ResponseError::new( + ErrorCode::REQUEST_CANCELLED, + "Formatting request was cancelled", + )) + }) + } + .boxed() + } +} diff --git a/src/server/request/initialize.rs b/src/server/request/initialize.rs new file mode 100644 index 0000000..f24c5dc --- /dev/null +++ b/src/server/request/initialize.rs @@ -0,0 +1,261 @@ +use std::path::PathBuf; + +use async_lsp::ResponseError; +use async_lsp::lsp_types::{ + CompletionOptions, DefinitionOptions, FileOperationFilter, FileOperationPattern, + FileOperationPatternKind, FileOperationRegistrationOptions, HoverOptions, + HoverProviderCapability, InitializeParams, InitializeResult, OneOf, RenameOptions, + ServerCapabilities, ServerInfo, TextDocumentSyncCapability, TextDocumentSyncKind, + WorkDoneProgressOptions, WorkspaceFileOperationsServerCapabilities, + WorkspaceFoldersServerCapabilities, WorkspaceServerCapabilities, WorkspaceSymbolOptions, +}; +use futures::FutureExt; +use futures::future::BoxFuture; +use serde_json::Value; + +use crate::ProtoLanguageServer; +use crate::config::client::ClientCapabilitiesSummary; +use crate::config::configure_workspace_folders; + +impl ProtoLanguageServer { + pub(in crate::server) fn initialize( + &mut self, + params: InitializeParams, + ) -> BoxFuture<'static, Result> { + log_client_info(¶ms); + + self.state.set_client_capabilities((¶ms).into()); + + let configs = self.configs.clone(); + let include_paths = parse_init_include_paths(¶ms); + let workspace = Some(build_workspace_capabilities(¶ms)); + let definition_provider = build_definition_provider(¶ms); + let hover_provider = build_hover_provider(¶ms); + let workspace_symbol_provider = build_workspace_symbol_provider(¶ms); + let rename_provider = build_rename_provider(¶ms); + + async move { + let folders = params.workspace_folders.as_deref().unwrap_or_default(); + let configured_folders = configure_workspace_folders(folders).await; + + { + let mut config = configs.write().await; + + if !include_paths.is_empty() { + tracing::info!( + "Setting include paths from initialization options: {:?}", + include_paths + ); + config.set_init_include_paths(include_paths); + } + + if params.workspace_folders.as_ref().is_none_or(Vec::is_empty) { + tracing::info!("Running in no workspace mode"); + config.no_workspace_mode(); + } else { + for folder in configured_folders { + tracing::info!("Workspace folder: {:?}", folder.uri); + config.add_workspace_folder(folder); + } + } + } + + let response = InitializeResult { + capabilities: ServerCapabilities { + // todo(): We might prefer incremental sync at some later stage + text_document_sync: Some(TextDocumentSyncCapability::Kind( + TextDocumentSyncKind::FULL, + )), + workspace, + definition_provider, + // hover_provider, + document_symbol_provider: Some(OneOf::Left(true)), + // workspace_symbol_provider, + completion_provider: Some(CompletionOptions::default()), + rename_provider, + document_formatting_provider: Some(OneOf::Left(true)), + document_range_formatting_provider: Some(OneOf::Left(true)), + references_provider: Some(OneOf::Left(true)), + ..ServerCapabilities::default() + }, + server_info: Some(ServerInfo { + name: env!("CARGO_PKG_NAME").to_string(), + version: Some(env!("CARGO_PKG_VERSION").to_string()), + }), + }; + + Ok(response) + } + .boxed() + } +} + +/// Parse `include_paths` from initialization options +#[inline] +fn parse_init_include_paths(params: &InitializeParams) -> Vec { + let Some(paths) = params + .initialization_options + .as_ref() + .and_then(|v| v["include_paths"].as_array()) + else { + return Vec::new(); + }; + + let trace_error = |path_value: &&Value| { + if !path_value.is_string() { + tracing::warn!( + "Invalid include path in initialization options: {:?}", + path_value + ); + } + }; + + paths + .iter() + .inspect(trace_error) + .filter_map(Value::as_str) + .map(PathBuf::from) + .collect() +} + +#[inline] +fn log_client_info(params: &InitializeParams) { + let name = params + .client_info + .as_ref() + .map_or("", |c| c.name.as_str()); + let version = params + .client_info + .as_ref() + .and_then(|c| c.version.as_deref()) + .unwrap_or(""); + + tracing::info!("Connected with client {name} {version}"); +} + +#[inline] +fn build_workspace_capabilities(params: &InitializeParams) -> WorkspaceServerCapabilities { + let file_operations = build_file_operations(params); + + WorkspaceServerCapabilities { + workspace_folders: Some(WorkspaceFoldersServerCapabilities { + supported: Some(true), + change_notifications: Some(OneOf::Left( + params + .capabilities + .workspace + .as_ref() + .and_then(|w| w.workspace_folders) + .unwrap_or_default(), + )), + }), + file_operations, + } +} + +#[inline] +fn build_file_operations( + params: &InitializeParams, +) -> Option { + let file_operations = params + .capabilities + .workspace + .as_ref()? + .file_operations + .as_ref()?; + + let did_create = file_operations.did_create.unwrap_or_default(); + let did_delete = file_operations.did_delete.unwrap_or_default(); + let did_rename = file_operations.did_rename.unwrap_or_default(); + + Some(WorkspaceFileOperationsServerCapabilities { + did_create: did_create.then_some(build_file_operation_filters()), + did_delete: did_delete.then_some(build_file_operation_filters()), + did_rename: did_rename.then_some(build_file_operation_filters()), + ..Default::default() + }) +} + +#[inline] +fn build_file_operation_filters() -> FileOperationRegistrationOptions { + FileOperationRegistrationOptions { + filters: vec![FileOperationFilter { + scheme: Some(String::from("file")), + pattern: FileOperationPattern { + glob: String::from("**/*.proto"), + matches: Some(FileOperationPatternKind::File), + ..Default::default() + }, + }], + } +} + +#[inline] +fn build_definition_provider(params: &InitializeParams) -> Option> { + params + .capabilities + .text_document + .as_ref() + .and_then(|cap| cap.definition.as_ref()) + .map(|_| { + OneOf::Right(DefinitionOptions { + work_done_progress_options: WorkDoneProgressOptions { + work_done_progress: Some(false), + }, + }) + }) +} + +#[inline] +fn build_hover_provider(params: &InitializeParams) -> Option { + params + .capabilities + .text_document + .as_ref() + .and_then(|cap| cap.hover.as_ref()) + .map(|_| { + HoverProviderCapability::Options(HoverOptions { + work_done_progress_options: WorkDoneProgressOptions { + work_done_progress: Some(false), + }, + }) + }) +} + +#[inline] +fn build_workspace_symbol_provider( + params: &InitializeParams, +) -> Option> { + let work_done_progress_options = + ClientCapabilitiesSummary::build_work_done_progress_options(params); + + params + .capabilities + .workspace + .as_ref() + .and_then(|cap| cap.symbol.as_ref()) + .map(|_| { + OneOf::Right(WorkspaceSymbolOptions { + resolve_provider: Some(false), + work_done_progress_options, + }) + }) +} + +#[inline] +fn build_rename_provider(params: &InitializeParams) -> Option> { + let work_done_progress_options = + ClientCapabilitiesSummary::build_work_done_progress_options(params); + + params + .capabilities + .text_document + .as_ref() + .and_then(|cap| cap.rename.as_ref()) + .map(|rc| { + OneOf::Right(RenameOptions { + prepare_provider: rc.prepare_support, + work_done_progress_options, + }) + }) +} diff --git a/src/server/request/shutdown.rs b/src/server/request/shutdown.rs new file mode 100644 index 0000000..ab2a3c7 --- /dev/null +++ b/src/server/request/shutdown.rs @@ -0,0 +1,14 @@ +use async_lsp::ResponseError; +use futures::future::{Ready, ready}; + +use crate::ProtoLanguageServer; + +impl ProtoLanguageServer { + pub(in crate::server) fn shutdown(&mut self, _params: ()) -> Ready> { + tracing::info!("Received shutdown request"); + + self.shutdown_token.cancel(); + + ready(Ok(())) + } +} diff --git a/src/server/state.rs b/src/server/state.rs new file mode 100644 index 0000000..c1b815f --- /dev/null +++ b/src/server/state.rs @@ -0,0 +1,115 @@ +use std::collections::HashMap; +use std::sync::atomic::Ordering; +use std::sync::{Arc, OnceLock, atomic::AtomicBool}; + +use async_lsp::lsp_types::Url; +use tokio::sync::RwLock; +use tree_sitter::{Query, QueryError}; + +use crate::config::client::ClientCapabilitiesSummary; +use crate::model::generate_metamodel_query; +use cache::CacheEntry; +use source::SourceEntry; + +pub(super) mod cache; +pub(super) mod error; +pub(super) mod finalizer; +// mod hover; +pub(super) mod index; +mod query; +// mod resolve; +pub(super) mod source; +pub(super) mod unique_uris; + +#[derive(Clone)] +pub(crate) struct ProtoLanguageState { + client_capabilities: Arc>, + sources: Arc>>, + documents: Arc>>, + metamodel_query: Arc, + is_indexing: Arc, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub struct DocumentVersion(pub i32); + +impl DocumentVersion { + pub const DISK: Self = DocumentVersion(-1); +} + +impl ProtoLanguageState { + pub fn new() -> Self { + let language: tree_sitter::Language = tree_sitter_proto::LANGUAGE.into(); + let trace_error = |e: &QueryError| { + tracing::error!( + "Critical SCM error: Failed to compile embedded Tree-sitter query for metadata extraction. Details: {:?}", + e + ); + }; + + let metamodel_query = Query::new(&language, &generate_metamodel_query()) + .inspect_err(trace_error) + .expect("Tree-sitter query compilation failed"); + + Self { + client_capabilities: Arc::default(), + sources: Arc::default(), + documents: Arc::default(), + metamodel_query: Arc::new(metamodel_query), + is_indexing: Arc::default(), + } + } + + pub async fn write_lock_all( + &self, + ) -> ( + tokio::sync::RwLockWriteGuard<'_, HashMap>, + tokio::sync::RwLockWriteGuard<'_, HashMap>, + ) { + let documents_guard = self.documents.write().await; + let sources_guard = self.sources.write().await; + + (documents_guard, sources_guard) + } + + #[inline] + pub fn set_client_capabilities(&self, capabilities: ClientCapabilitiesSummary) { + self.client_capabilities + .set(capabilities) + .expect("Attempted to initialize client capabilities more than once per session!"); + } + + #[inline] + pub fn client_capabilities(&self) -> ClientCapabilitiesSummary { + self.client_capabilities.get().copied().unwrap_or_default() + } + + #[inline] + pub fn set_indexing(&self, value: bool) { + self.is_indexing.store(value, Ordering::Relaxed); + } + + #[inline] + pub fn is_indexing(&self) -> bool { + self.is_indexing.load(Ordering::Relaxed) + } + + #[cfg(test)] + pub async fn add_dirty_file_for_test( + &self, + uri: &Url, + content: &str, + version: DocumentVersion, + ) { + let content = Arc::new(content.to_owned()); + + { + use crate::server::state::cache::DocumentsGuardExt; + + let (mut documents, mut sources) = self.write_lock_all().await; + + documents.set_dirty(uri); + sources.insert(uri.clone(), SourceEntry { version, content }); + } + } +} diff --git a/src/server/state/cache.rs b/src/server/state/cache.rs new file mode 100644 index 0000000..2dbac33 --- /dev/null +++ b/src/server/state/cache.rs @@ -0,0 +1,122 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use async_lsp::lsp_types::Url; +use futures::future::{BoxFuture, Shared}; +use tokio::sync::RwLockWriteGuard; +use tokio_util::sync::CancellationToken; + +use super::DocumentVersion; +use super::error::CoreError; +use crate::document::ProtoDocument; + +pub(super) type ParsingResult = Result, CoreError>; + +pub(super) type SharedParsingFuture = Shared>; + +#[derive(Clone)] +pub enum CacheEntry { + Ready { + version: DocumentVersion, + result: ParsingResult, + }, + Pending { + version: DocumentVersion, + future: SharedParsingFuture, + cancel_token: CancellationToken, + }, + Dirty, +} + +impl CacheEntry { + pub(super) fn is_transient_error(result: &Result, CoreError>) -> bool { + match result { + Ok(_) + | Err( + CoreError::DiskIoError { .. } + | CoreError::InvalidPath { .. } + | CoreError::InvalidExtension { .. } + | CoreError::SourceNotFound { .. } + | CoreError::InternalError { .. }, + ) => false, + Err( + CoreError::DocumentNotFound { .. } + | CoreError::ParserError { .. } + | CoreError::RequestCancelled { .. } + | CoreError::FileTooLarge { .. }, + ) => true, + } + } + + pub(super) fn finalize( + &mut self, + version: DocumentVersion, + parse_result: Result, CoreError>, + ) { + if let Self::Pending { + version: current_version, + .. + } = self + && version == *current_version + { + *self = Self::Ready { + version, + result: parse_result, + }; + } + } + + #[inline] + pub fn cancel(&self) { + if let Self::Pending { cancel_token, .. } = self { + cancel_token.cancel(); + } + } +} + +pub(super) fn get_uri(result: &Result, CoreError>) -> &Url { + match result { + Ok(d) => &d.uri, + Err( + CoreError::DiskIoError { uri, .. } + | CoreError::DocumentNotFound { uri } + | CoreError::InternalError { uri } + | CoreError::InvalidPath { uri } + | CoreError::InvalidExtension { uri } + | CoreError::FileTooLarge { uri, .. } + | CoreError::ParserError { uri } + | CoreError::RequestCancelled { uri } + | CoreError::SourceNotFound { uri }, + ) => uri, + } +} +pub trait DocumentsGuardExt { + /// Forces the document state to `Dirty`, cancels any pending background tasks, + /// and returns `true` if the document already existed in the cache. + fn set_dirty(&mut self, uri: &Url) -> bool; + + /// Completely removes the document from the cache, triggers `cancel()` if it was `Pending`, + /// and returns `Some(CacheEntry)` if the document existed. + fn cancel_and_remove(&mut self, uri: &Url) -> Option; +} + +impl DocumentsGuardExt for RwLockWriteGuard<'_, HashMap> { + fn set_dirty(&mut self, uri: &Url) -> bool { + if let Some(entry) = self.get_mut(uri) { + entry.cancel(); + *entry = CacheEntry::Dirty; + + return true; + } + + self.insert(uri.clone(), CacheEntry::Dirty).is_some() + } + + fn cancel_and_remove(&mut self, uri: &Url) -> Option { + let entry = self.remove(uri)?; + + entry.cancel(); + + Some(entry) + } +} diff --git a/src/server/state/error.rs b/src/server/state/error.rs new file mode 100644 index 0000000..d842268 --- /dev/null +++ b/src/server/state/error.rs @@ -0,0 +1,68 @@ +use async_lsp::lsp_types::Url; +use async_lsp::{ErrorCode, ResponseError}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum CoreError { + DocumentNotFound { uri: Url }, + SourceNotFound { uri: Url }, + InvalidPath { uri: Url }, + InvalidExtension { uri: Url }, + FileTooLarge { uri: Url, size: u64 }, + DiskIoError { uri: Url, details: String }, + ParserError { uri: Url }, + RequestCancelled { uri: Url }, + InternalError { uri: Url }, +} + +impl From for ResponseError { + fn from(err: CoreError) -> Self { + match err { + CoreError::DocumentNotFound { uri } => ResponseError::new( + ErrorCode::INVALID_PARAMS, + format!("Protols cache failure: Document metadata wrapper not found {uri}"), + ), + + CoreError::SourceNotFound { uri } => ResponseError::new( + ErrorCode::INTERNAL_ERROR, + format!("Protols cache corruption: Raw source text missing for tracked file {uri}"), + ), + + CoreError::InvalidPath { uri } => ResponseError::new( + ErrorCode::INTERNAL_ERROR, + format!("Protols failed to convert URI to OS path: {uri}"), + ), + + CoreError::InvalidExtension { uri } => ResponseError::new( + ErrorCode::REQUEST_FAILED, + format!("Unsupported file extension. Protols only handles '.proto' files: {uri}"), + ), + + CoreError::FileTooLarge { uri, size } => ResponseError::new( + ErrorCode::REQUEST_FAILED, + format!("File exceeds maximum allowed size for parsing ({size} bytes): {uri}"), + ), + + CoreError::DiskIoError { details, uri } => ResponseError::new( + ErrorCode::INTERNAL_ERROR, + format!("Protols failed to read from disk: {uri}. Details: {details}"), + ), + + CoreError::ParserError { uri } => ResponseError::new( + ErrorCode::INTERNAL_ERROR, + format!("Protols failed to create parser for URI {uri}"), + ), + + CoreError::RequestCancelled { uri } => ResponseError::new( + ErrorCode::REQUEST_CANCELLED, + format!("LSP request was aborted or superseded by a newer action for URI {uri}"), + ), + + CoreError::InternalError { uri } => ResponseError::new( + ErrorCode::INTERNAL_ERROR, + format!( + "Protols internal panic: Batch processing engine lost requested file map entry {uri}" + ), + ), + } + } +} diff --git a/src/server/state/finalizer.rs b/src/server/state/finalizer.rs new file mode 100644 index 0000000..7bd4b2c --- /dev/null +++ b/src/server/state/finalizer.rs @@ -0,0 +1,28 @@ +use std::pin::Pin; +use std::task::{Context, Poll}; + +pub(crate) struct FinalizerToken { + inner: tokio::task::JoinHandle<()>, +} + +impl FinalizerToken { + /// Creates a dummy completed finalizer token that resolves instantly. + pub(crate) fn dummy() -> Self { + Self { + inner: tokio::spawn(async {}), + } + } + + pub(crate) fn from_handle(task_handle: tokio::task::JoinHandle<()>) -> Self { + Self { inner: task_handle } + } +} + +impl Future for FinalizerToken { + type Output = Result<(), tokio::task::JoinError>; + + #[inline] + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + Pin::new(&mut self.inner).poll(cx) + } +} diff --git a/src/server/state/index.rs b/src/server/state/index.rs new file mode 100644 index 0000000..14541be --- /dev/null +++ b/src/server/state/index.rs @@ -0,0 +1,73 @@ +use async_lsp::lsp_types::Url; +use tokio_util::sync::CancellationToken; +use walkdir::{DirEntry, WalkDir}; + +use crate::server::progress::{LspProgressReporter, OptionReporterExt}; +use crate::server::state::unique_uris::UniqueUris; +use crate::utils::OutermostPaths; + +impl super::ProtoLanguageState { + pub async fn index_folders( + &self, + paths: OutermostPaths, + cancel_token: CancellationToken, + progress_reporter: Option, + ) { + if paths.is_empty() { + return; + } + + let cancel_token_clone = cancel_token.clone(); + + let uris: Vec<_> = tokio::task::spawn_blocking(move || { + let trace_walkdir_error = |item: &Result| { + if let Err(err) = item { + tracing::warn!("Failed to access path during .proto files scan: {err}"); + } + }; + + paths + .into_iter() + .take_while(|_| !cancel_token.is_cancelled()) + .filter(|p| p.is_dir()) + .flat_map(|dir| { + WalkDir::new(dir) + .follow_links(false) + .into_iter() + .filter_entry(|e| !cancel_token.is_cancelled() && !is_hidden(e)) + .inspect(trace_walkdir_error) + .filter_map(Result::ok) + }) + .take_while(|_| !cancel_token.is_cancelled()) + .filter(|de| de.path().extension().is_some_and(|ext| ext == "proto")) + .filter_map(|de| Url::from_file_path(de.path()).ok()) + .collect() + }) + .await + .unwrap_or_default(); + + progress_reporter.report(20, "Scan completed. Starting file parsing..."); + + if !uris.is_empty() && !cancel_token_clone.is_cancelled() { + tracing::info!( + "Workspace scan completed. Booking cache for {} .proto files...", + uris.len() + ); + + let (finalized, _) = self + .query_documents_batch( + UniqueUris::new(uris), + cancel_token_clone, + progress_reporter.with_range(20, 100), + false, + ) + .await; + + let _ = finalized.await; + } + } +} + +fn is_hidden(entry: &walkdir::DirEntry) -> bool { + entry.file_name().as_encoded_bytes().starts_with(b".") +} diff --git a/src/server/state/query.rs b/src/server/state/query.rs new file mode 100644 index 0000000..0a67002 --- /dev/null +++ b/src/server/state/query.rs @@ -0,0 +1,362 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use async_lsp::lsp_types::Url; +use futures::future::Either; +use futures::stream::{FuturesUnordered, StreamExt}; +use futures::{FutureExt, Stream}; +use tokio_util::sync::CancellationToken; +use tree_sitter::LanguageError; + +use super::cache::get_uri; +use super::cache::{ParsingResult, SharedParsingFuture}; +use super::error::CoreError; +use super::finalizer::FinalizerToken; +use super::unique_uris::UniqueUris; +use super::{CacheEntry, DocumentVersion, ProtoLanguageState}; +use crate::document::ProtoDocument; +use crate::server::progress::{LspProgressReporter, OptionReporterExt}; + +const WRITE_CHUNK_SIZE: usize = 16; +const MAX_PROTO_FILE_SIZE_BYTES: u64 = 1024 * 1024; + +enum TaskSource { + FromMemory(Arc), + FromDisk, +} + +type ParsingTask = ( + Url, + TaskSource, + tokio::sync::oneshot::Sender, CoreError>>, + CancellationToken, +); + +impl ProtoLanguageState { + /// Processes a batch of documents. + /// + /// # Cancellation safety + /// + /// This method is cancel safe. + pub async fn query_documents_batch( + &self, + uris: UniqueUris, + batch_cancel_token: CancellationToken, + progress_reporter: Option, + need_results: bool, + ) -> (FinalizerToken, impl Stream) { + { + let documents = self.documents.read().await; + if let Some((cache_hits, client_futures)) = + Self::probe_cache_fast_path(&uris, need_results, &documents) + { + return ( + FinalizerToken::dummy(), + Either::Left(futures::stream::iter(cache_hits).chain(client_futures)), + ); + } + } + + let mut cache_hits = Vec::new(); + let mut client_futures = FuturesUnordered::new(); + let finalizer_futures = FuturesUnordered::new(); + let mut tasks = Vec::new(); + + { + let mut documents = self.documents.write().await; + let sources = self.sources.read().await; + + for uri in uris { + let entry = documents.get(&uri); + + match entry { + Some(CacheEntry::Ready { result, .. }) + if !CacheEntry::is_transient_error(result) => + { + cache_hits.extend(need_results.then(|| result.clone())); + continue; + } + Some(CacheEntry::Pending { future, .. }) => { + client_futures.extend(need_results.then(|| future.clone())); + continue; + } + _ => {} + } + + let source = sources.get(&uri); + + if source.is_none() + && matches!(entry, Some(CacheEntry::Dirty | CacheEntry::Ready { .. })) + { + let result = Err(CoreError::SourceNotFound { uri: uri.clone() }); + + cache_hits.extend(need_results.then(|| result.clone())); + documents.insert( + uri.clone(), + CacheEntry::Ready { + version: DocumentVersion::DISK, + result, + }, + ); + + continue; + } + + let (expected_version, task_source) = source.map_or_else( + || (DocumentVersion::DISK, TaskSource::FromDisk), + |se| (se.version, TaskSource::FromMemory(se.content.clone())), + ); + + let (tx, rx) = tokio::sync::oneshot::channel(); + let cancel_token = batch_cancel_token.child_token(); + let cancel_token_clone = cancel_token.clone(); + let uri_clone = uri.clone(); + + let future = async move { + let uri = uri_clone; + cancel_token_clone + .run_until_cancelled(rx) + .await + .map(|res| { + res.unwrap_or(Err(CoreError::RequestCancelled { uri: uri.clone() })) + }) + .unwrap_or(Err(CoreError::RequestCancelled { uri })) + } + .boxed() + .shared(); + + let version = expected_version; + + finalizer_futures + .push(future.clone().map(move |result| (expected_version, result))); + + client_futures.extend(need_results.then(|| future.clone())); + + documents.insert( + uri.clone(), + CacheEntry::Pending { + version, + future, + cancel_token: cancel_token.clone(), + }, + ); + + tasks.push((uri, task_source, tx, cancel_token)); + } + } + + self.run_blocking_parser(tasks); + + ( + self.run_cache_finalizer(finalizer_futures, progress_reporter), + Either::Right(futures::stream::iter(cache_hits).chain(client_futures)), + ) + } + + fn probe_cache_fast_path( + uris: &UniqueUris, + need_results: bool, + documents: &HashMap, + ) -> Option<(Vec, FuturesUnordered)> { + let mut cache_hits = Vec::new(); + let mut client_futures = FuturesUnordered::new(); + + for uri in uris { + match documents.get(uri) { + Some(CacheEntry::Ready { result, .. }) + if !CacheEntry::is_transient_error(result) => + { + cache_hits.extend(need_results.then(|| result.clone())); + } + Some(CacheEntry::Pending { future, .. }) => { + client_futures.extend(need_results.then(|| future.clone())); + } + _ => { + return None; + } + } + } + + Some((cache_hits, client_futures)) + } + + fn run_blocking_parser(&self, tasks: Vec) { + if tasks.is_empty() { + return; + } + + let query = self.metamodel_query.clone(); + tokio::task::spawn_blocking(move || { + let mut parser = tree_sitter::Parser::new(); + + let trace_error = |error: &LanguageError| { + tracing::error!( + ?error, + "Critical initialization failure: Failed to set Tree-sitter Protobuf language parser" + ); + }; + + parser + .set_language(&tree_sitter_proto::LANGUAGE.into()) + .inspect_err(trace_error) + .expect("Tree-sitter parser: setting language failed"); + + for (uri, task_source, tx, cancel_token) in tasks { + if cancel_token.is_cancelled() { + let _ = tx.send(Err(CoreError::RequestCancelled { uri })); + continue; + } + + let mut disk_buffer = Vec::new(); + + if let TaskSource::FromDisk = &task_source { + let Ok(path) = uri.to_file_path() else { + let _ = tx.send(Err(CoreError::InvalidPath { uri })); + continue; + }; + + if path.extension().is_none_or(|ext| ext != "proto") { + let _ = tx.send(Err(CoreError::InvalidExtension { uri })); + continue; + } + + match std::fs::metadata(&path) { + Err(e) => { + let _ = tx.send(Err(CoreError::DiskIoError { + uri, + details: e.to_string(), + })); + continue; + } + Ok(metadata) => { + let size = metadata.len(); + + if size > MAX_PROTO_FILE_SIZE_BYTES { + tracing::warn!( + %uri, + size, + "Refusing to parse file: size exceeds maximum limit of {} bytes", + MAX_PROTO_FILE_SIZE_BYTES + ); + let _ = tx.send(Err(CoreError::FileTooLarge { uri, size })); + continue; + } + + if !metadata.is_file() { + let _ = tx.send(Err(CoreError::InvalidPath { uri })); + continue; + } + } + } + + match std::fs::read(&path) { + Ok(bytes) => { + disk_buffer = bytes; + } + Err(error) => { + let _ = tx.send(Err(CoreError::DiskIoError { + uri, + details: error.to_string(), + })); + continue; + } + } + } + + let bytes_ref: &[u8] = match &task_source { + TaskSource::FromMemory(content) => content.as_bytes(), + TaskSource::FromDisk => &disk_buffer, + }; + + if cancel_token.is_cancelled() { + let _ = tx.send(Err(CoreError::RequestCancelled { uri })); + continue; + } + + let result = + ProtoDocument::try_from_input(uri.clone(), bytes_ref, &query, &mut parser) + .map(Arc::new) + .ok_or(CoreError::ParserError { uri }); + + let _ = tx.send(result); + } + }); + } + + fn run_cache_finalizer( + &self, + futures: FuturesUnordered, + progress_reporter: Option, + ) -> FinalizerToken + where + Fut: Future + Send + 'static, + { + if futures.is_empty() { + tracing::warn!( + "Document cache finalizer called with 0 tasks. This might indicate lost futures." + ); + + return FinalizerToken::dummy(); + } + + let total = u32::try_from(futures.len()).unwrap_or(u32::MAX); + let mut completed = 0; + let mut chunks = futures.ready_chunks(WRITE_CHUNK_SIZE); + let documents_clone = self.documents.clone(); + + let handle = tokio::spawn(async move { + while let Some(chunk) = chunks.next().await { + { + let mut documents = documents_clone.write().await; + + for (expected_version, result) in chunk { + if let Some(document) = documents.get_mut(get_uri(&result)) { + document.finalize(expected_version, result); + } + + completed += 1; + } + } + + let percentage = completed * 100 / total; + progress_reporter.report(percentage, &format!("Parsing [{completed}/{total}]")); + + tokio::task::yield_now().await; + } + }); + + FinalizerToken::from_handle(handle) + } + + pub async fn get_document(&self, uri: &Url, cancel_token: CancellationToken) -> ParsingResult { + let uris = UniqueUris::new(uri.clone()); + + let (_, mut stream) = self + .query_documents_batch(uris, cancel_token, None, true) + .await; + + match stream.next().await { + Some(result) if get_uri(&result) == uri => result, + + Some(result) => { + let returned_uri = get_uri(&result); + tracing::error!( + ?returned_uri, + expected_uri = ?uri, + "Bug: query_documents_batch returned a mismatched URL" + ); + Err(CoreError::InternalError { uri: uri.clone() }) + } + + None => Err(CoreError::DocumentNotFound { uri: uri.clone() }), + } + } + + pub async fn get_document_text(&self, uri: &Url) -> Option> { + self.sources + .read() + .await + .get(uri) + .map(|se| se.content.clone()) + } +} diff --git a/src/server/state/source.rs b/src/server/state/source.rs new file mode 100644 index 0000000..1a1546c --- /dev/null +++ b/src/server/state/source.rs @@ -0,0 +1,9 @@ +use std::sync::Arc; + +use super::DocumentVersion; + +#[derive(Clone)] +pub struct SourceEntry { + pub version: DocumentVersion, + pub content: Arc, +} diff --git a/src/server/state/unique_uris.rs b/src/server/state/unique_uris.rs new file mode 100644 index 0000000..53b5ca9 --- /dev/null +++ b/src/server/state/unique_uris.rs @@ -0,0 +1,95 @@ +use std::{collections::HashSet, ops::Deref}; + +use async_lsp::lsp_types::Url; + +/// A type-safe container that always holds unique and sorted URLs. +pub struct UniqueUris { + uris: Vec, +} + +impl UniqueUris { + #[inline] + pub fn new(source: T) -> Self + where + T: IntoUniqueUris, + { + source.into_unique_uris() + } + + #[inline] + pub fn iter(&self) -> std::slice::Iter<'_, Url> { + self.uris.iter() + } + + #[inline] + pub fn into_inner(self) -> Vec { + self.uris + } +} + +impl Deref for UniqueUris { + type Target = [Url]; + + #[inline] + fn deref(&self) -> &Self::Target { + &self.uris + } +} + +impl IntoIterator for UniqueUris { + type Item = Url; + type IntoIter = std::vec::IntoIter; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.uris.into_iter() + } +} + +impl<'a> IntoIterator for &'a UniqueUris { + type Item = &'a Url; + type IntoIter = std::slice::Iter<'a, Url>; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.uris.iter() + } +} + +pub trait IntoUniqueUris { + fn into_unique_uris(self) -> UniqueUris; +} + +impl IntoUniqueUris for HashSet { + #[inline] + fn into_unique_uris(self) -> UniqueUris { + let uris: Vec = self.into_iter().collect(); + UniqueUris { uris } + } +} + +impl IntoUniqueUris for std::collections::hash_map::IntoKeys { + #[inline] + fn into_unique_uris(self) -> UniqueUris { + UniqueUris { + uris: self.collect(), + } + } +} + +impl IntoUniqueUris for Vec { + #[inline] + fn into_unique_uris(self) -> UniqueUris { + let mut uris = self; + uris.sort_unstable(); + uris.dedup(); + UniqueUris { uris } + } +} + +impl IntoUniqueUris for Url { + #[inline] + fn into_unique_uris(self) -> UniqueUris { + UniqueUris { uris: vec![self] } + } +} diff --git a/src/state/definition.rs b/src/state/definition.rs index 0451b26..b41902b 100644 --- a/src/state/definition.rs +++ b/src/state/definition.rs @@ -1,6 +1,7 @@ use std::path::PathBuf; use async_lsp::lsp_types::{Location, Position, Range, Url}; +use tokio_util::sync::CancellationToken; use crate::{ model::{ElementKind, SpatialEntry}, @@ -15,8 +16,14 @@ impl ProtoLanguageState { /// cursor: an `import` statement jumps to the imported file, while any /// other symbol (or a type reference) is resolved to its declaration via /// the shared cross-file name resolution engine. - pub fn definition(&self, uri: &Url, pos: Position, ipath: &[PathBuf]) -> Vec { - let Some(document) = self.get_document(uri) else { + pub async fn definition( + &self, + uri: &Url, + pos: Position, + ipath: &[PathBuf], + cancel_token: CancellationToken, + ) -> Vec { + let Ok(document) = self.get_document(uri, cancel_token.clone()).await else { return vec![]; }; let Some(SpatialEntry { element_id, .. }) = document.find_entry_at_position(pos) else { @@ -39,7 +46,7 @@ impl ProtoLanguageState { }]; } - let Some(fqn) = self.resolve_target_fqn(uri, pos) else { + let Some(fqn) = self.resolve_target_fqn(uri, pos, cancel_token).await else { return vec![]; }; self.declarations_for_fqn(&fqn) @@ -50,6 +57,7 @@ impl ProtoLanguageState { mod test { use async_lsp::lsp_types::{Position, Url}; use std::path::PathBuf; + use tokio_util::sync::CancellationToken; use insta::assert_yaml_snapshot; @@ -105,28 +113,38 @@ mod test { assert_yaml_snapshot!(state.resolve_identifier_locations("com.utility", "Baz")); } - #[test] - fn test_definition_position_based() { + #[tokio::test(flavor = "current_thread")] + async fn test_definition_position_based() { let (ipath, a_uri, b_uri, _c_uri, state) = setup_workspace(); // Cursor on the `Author` message declaration name in b.proto. - assert_yaml_snapshot!(state.definition( - &b_uri, - Position { - line: 5, - character: 10 - }, - &ipath - )); + assert_yaml_snapshot!( + state + .definition( + &b_uri, + Position { + line: 5, + character: 10 + }, + &ipath, + CancellationToken::new() + ) + .await + ); // Cursor on the `Author` type reference inside a field in a.proto. - assert_yaml_snapshot!(state.definition( - &a_uri, - Position { - line: 11, - character: 5 - }, - &ipath - )); + assert_yaml_snapshot!( + state + .definition( + &a_uri, + Position { + line: 11, + character: 5 + }, + &ipath, + CancellationToken::new() + ) + .await + ); // Cursor on empty whitespace -> no definition. assert!( state @@ -136,14 +154,16 @@ mod test { line: 0, character: 0 }, - &ipath + &ipath, + CancellationToken::new() ) + .await .is_empty() ); } - #[test] - fn test_definition_import_position_based() { + #[tokio::test(flavor = "current_thread")] + async fn test_definition_import_position_based() { let ipath = vec![std::env::current_dir().unwrap().join("src/state/input")]; let a_uri = "file://input/a.proto".parse().unwrap(); let mut state: ProtoLanguageState = ProtoLanguageState::new(); @@ -157,17 +177,20 @@ mod test { ); // Cursor on the `import "c.proto"` statement -> jump to the file. - let loc = state.definition( - &a_uri, - Position { - line: 4, - character: 10, - }, - &ipath, - ); + let loc = state + .definition( + &a_uri, + Position { + line: 4, + character: 10, + }, + &ipath, + CancellationToken::new(), + ) + .await; assert_yaml_snapshot!(loc, {"[0].uri" => insta::dynamic_redaction(|c, _| { assert!(c.as_str().unwrap().ends_with("c.proto")); - "file:///c.proto".to_string() + "file:///c.proto".to_owned() })}); } diff --git a/src/state/hover.rs b/src/state/hover.rs index ab73376..91a5f6f 100644 --- a/src/state/hover.rs +++ b/src/state/hover.rs @@ -1,7 +1,8 @@ use async_lsp::lsp_types::{Hover, HoverContents, MarkupContent, MarkupKind, Position, Url}; +use tokio_util::sync::CancellationToken; use crate::model::{ElementKind, ModelElement, SpatialEntry}; -use crate::state::ProtoLanguageState; +use crate::state::{CoreError, ProtoLanguageState}; use crate::utils::is_position_inside_range; impl ProtoLanguageState { @@ -12,8 +13,13 @@ impl ProtoLanguageState { /// /// Returns `Some(Hover)` containing the unified markdown payload, or `None` /// if the token carries no hoverable metadata. - pub fn hover(&self, uri: &Url, position: Position) -> Option { - let current_document = self.get_document(uri)?; + pub async fn hover( + &self, + uri: &Url, + position: Position, + cancel_token: CancellationToken, + ) -> Result, CoreError> { + let current_document = self.get_document(uri, cancel_token).await?; let SpatialEntry { element_id, range } = current_document.find_entry_at_position(position).copied()?; @@ -86,11 +92,13 @@ impl ModelElement { mod test { use async_lsp::lsp_types::Position; use insta::assert_yaml_snapshot; + use tokio_util::sync::CancellationToken; use crate::config::Config; use crate::state::ProtoLanguageState; - #[test] - fn workspace_test_hover() { + + #[tokio::test(flavor = "current_thread")] + async fn workspace_test_hover() { let ipath = vec![std::env::current_dir().unwrap().join("src/state/input")]; let a_uri = "file://input/a.proto".parse().unwrap(); let b_uri = "file://input/b.proto".parse().unwrap(); @@ -108,74 +116,128 @@ mod test { state.upsert_file(&c_uri, c, &ipath, 2, &Config::default(), false); state.upsert_file(&x_uri, x, &ipath, 2, &Config::default(), false); - assert_yaml_snapshot!(state.hover( - &a_uri, - Position { - line: 15, - character: 10 - } - )); - assert_yaml_snapshot!(state.hover( - &a_uri, - Position { - line: 11, - character: 6 - } - )); - assert_yaml_snapshot!(state.hover( - &b_uri, - Position { - line: 10, - character: 7 - } - )); - assert_yaml_snapshot!(state.hover( - &a_uri, - Position { - line: 12, - character: 14 - } - )); - assert_yaml_snapshot!(state.hover( - &a_uri, - Position { - line: 13, - character: 16 - } - )); - assert_yaml_snapshot!(state.hover( - &c_uri, - Position { - line: 12, - character: 5 - } - )); - assert_yaml_snapshot!(state.hover( - &a_uri, - Position { - line: 14, - character: 10 - } - )); - assert_yaml_snapshot!(state.hover( - &x_uri, - Position { - line: 9, - character: 18 - } - )); + assert_yaml_snapshot!( + state + .hover( + &a_uri, + Position { + line: 15, + character: 10 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); + assert_yaml_snapshot!( + state + .hover( + &a_uri, + Position { + line: 11, + character: 6 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); + assert_yaml_snapshot!( + state + .hover( + &b_uri, + Position { + line: 10, + character: 7 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); + assert_yaml_snapshot!( + state + .hover( + &a_uri, + Position { + line: 12, + character: 14 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); + assert_yaml_snapshot!( + state + .hover( + &a_uri, + Position { + line: 13, + character: 16 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); + assert_yaml_snapshot!( + state + .hover( + &c_uri, + Position { + line: 12, + character: 5 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); + assert_yaml_snapshot!( + state + .hover( + &a_uri, + Position { + line: 14, + character: 10 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); + assert_yaml_snapshot!( + state + .hover( + &x_uri, + Position { + line: 9, + character: 18 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); // relative path hover - assert_yaml_snapshot!(state.hover( - &x_uri, - Position { - line: 10, - character: 4 - } - )); + assert_yaml_snapshot!( + state + .hover( + &x_uri, + Position { + line: 10, + character: 4 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); } - #[test] - fn test_hover_builtin_and_wellknown() { + #[tokio::test(flavor = "current_thread")] + async fn test_hover_builtin_and_wellknown() { let ipath = vec![]; let uri = "file:///hover.proto".parse().unwrap(); let mut state: ProtoLanguageState = ProtoLanguageState::new(); @@ -196,28 +258,46 @@ mod test { ); // Hover over the builtin `string` type. - assert_yaml_snapshot!(state.hover( - &uri, - Position { - line: 3, - character: 3 - } - )); + assert_yaml_snapshot!( + state + .hover( + &uri, + Position { + line: 3, + character: 3 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); // Hover over the field name `title`. - assert_yaml_snapshot!(state.hover( - &uri, - Position { - line: 3, - character: 11 - } - )); + assert_yaml_snapshot!( + state + .hover( + &uri, + Position { + line: 3, + character: 11 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); // Hover over the well-known `google.protobuf.Any` type. - assert_yaml_snapshot!(state.hover( - &uri, - Position { - line: 4, - character: 3 - } - )); + assert_yaml_snapshot!( + state + .hover( + &uri, + Position { + line: 4, + character: 3 + }, + CancellationToken::new() + ) + .await + .unwrap() + ); } } diff --git a/src/state/index.rs b/src/state/index.rs new file mode 100644 index 0000000..8e318d1 --- /dev/null +++ b/src/state/index.rs @@ -0,0 +1,75 @@ +use async_lsp::lsp_types::Url; +use tokio_util::sync::CancellationToken; +use walkdir::{DirEntry, WalkDir}; + +use crate::{ + server::progress::{LspProgressReporter, OptionReporterExt}, + state::unique_uris::UniqueUris, + utils::OutermostPaths, +}; + +impl super::ProtoLanguageState { + pub async fn index_folders( + &self, + paths: OutermostPaths, + cancel_token: CancellationToken, + progress_reporter: Option, + ) { + if paths.is_empty() { + return; + } + + let cancel_token_clone = cancel_token.clone(); + + let uris: Vec<_> = tokio::task::spawn_blocking(move || { + let trace_walkdir_error = |item: &Result| { + if let Err(err) = item { + tracing::warn!("Failed to access path during .proto files scan: {err}"); + } + }; + + paths + .into_iter() + .take_while(|_| !cancel_token.is_cancelled()) + .filter(|p| p.is_dir()) + .flat_map(|dir| { + WalkDir::new(dir) + .follow_links(false) + .into_iter() + .filter_entry(|e| !cancel_token.is_cancelled() && !is_hidden(e)) + .inspect(trace_walkdir_error) + .filter_map(Result::ok) + }) + .take_while(|_| !cancel_token.is_cancelled()) + .filter(|de| de.path().extension().is_some_and(|ext| ext == "proto")) + .filter_map(|de| Url::from_file_path(de.path()).ok()) + .collect() + }) + .await + .unwrap_or_default(); + + progress_reporter.report(20, "Scan completed. Starting file parsing..."); + + if !uris.is_empty() && !cancel_token_clone.is_cancelled() { + tracing::info!( + "Workspace scan completed. Booking cache for {} .proto files...", + uris.len() + ); + + let (finalized, _) = self + .query_documents_batch( + UniqueUris::new(uris), + cancel_token_clone, + progress_reporter.with_range(20, 100), + false, + ) + .await; + + finalized.await; + } + } +} + +fn is_hidden(entry: &walkdir::DirEntry) -> bool { + entry.file_name().as_encoded_bytes().starts_with(b".") +} diff --git a/src/state/mod.rs b/src/state/mod.rs index 1fc9c40..4bfca23 100644 --- a/src/state/mod.rs +++ b/src/state/mod.rs @@ -1,23 +1,37 @@ mod definition; mod hover; +mod index; mod rename; mod resolve; +pub(super) mod unique_uris; mod workspace_symbol; +use futures::Stream; +use futures::future::{BoxFuture, Either, FutureExt, Shared}; +use futures::stream::{FuturesUnordered, StreamExt}; +use std::pin::Pin; +use std::sync::OnceLock; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::task::{Context, Poll}; use std::{ collections::{HashMap, HashSet}, path::{Path, PathBuf}, - sync::{Arc, Mutex, RwLock, mpsc::Sender}, + sync::{Arc, Mutex, mpsc::Sender}, }; +use tokio_util::sync::CancellationToken; use tracing::info; use async_lsp::lsp_types::{ - CompletionItem, CompletionItemKind, Location, OneOf, ProgressParamsValue, + CompletionItem, CompletionItemKind, FileDelete, Location, OneOf, ProgressParamsValue, PublishDiagnosticsParams, Range, SymbolKind, SymbolTag, Url, WorkspaceSymbol, }; -use tree_sitter::{Query, QueryError}; +use tokio::sync::{RwLock, RwLockWriteGuard, oneshot}; +use tree_sitter::{LanguageError, Query, QueryError}; use walkdir::WalkDir; +use crate::config::client::ClientCapabilitiesSummary; +use crate::server::progress::{LspProgressReporter, OptionReporterExt}; +use crate::state::unique_uris::UniqueUris; use crate::{ config::Config, document::{ProtoDocument, ProtoParser}, @@ -25,12 +39,224 @@ use crate::{ protoc::collect_diagnostics, }; +const MAX_PROTO_FILE_SIZE_BYTES: u64 = 1024 * 1024; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CoreError { + DocumentNotFound { uri: Url }, + SourceNotFound { uri: Url }, + InvalidPath { uri: Url }, + InvalidExtension { uri: Url }, + FileTooLarge { uri: Url, size: u64 }, + DiskIoError { uri: Url, details: String }, + ParserError { uri: Url }, + RequestCancelled { uri: Url }, + InternalError { uri: Url }, +} + +impl From for async_lsp::ResponseError { + fn from(err: CoreError) -> Self { + match err { + CoreError::DocumentNotFound { uri } => async_lsp::ResponseError::new( + async_lsp::ErrorCode::INVALID_PARAMS, + format!("Protols cache failure: Document metadata wrapper not found {uri}"), + ), + + CoreError::SourceNotFound { uri } => async_lsp::ResponseError::new( + async_lsp::ErrorCode::INTERNAL_ERROR, + format!("Protols cache corruption: Raw source text missing for tracked file {uri}"), + ), + + CoreError::InvalidPath { uri } => async_lsp::ResponseError::new( + async_lsp::ErrorCode::INTERNAL_ERROR, + format!("Protols failed to convert URI to OS path: {uri}"), + ), + + CoreError::InvalidExtension { uri } => async_lsp::ResponseError::new( + async_lsp::ErrorCode::REQUEST_FAILED, + format!("Unsupported file extension. Protols only handles '.proto' files: {uri}"), + ), + + CoreError::FileTooLarge { uri, size } => async_lsp::ResponseError::new( + async_lsp::ErrorCode::REQUEST_FAILED, + format!( + "File exceeds maximum allowed size for parsing ({} bytes): {uri}", + size + ), + ), + + CoreError::DiskIoError { details, uri } => async_lsp::ResponseError::new( + async_lsp::ErrorCode::INTERNAL_ERROR, + format!("Protols failed to read from disk: {uri}. Details: {details}"), + ), + + CoreError::ParserError { uri } => async_lsp::ResponseError::new( + async_lsp::ErrorCode::INTERNAL_ERROR, + format!("Protols failed to create parser for URI {uri}"), + ), + + CoreError::RequestCancelled { uri } => async_lsp::ResponseError::new( + async_lsp::ErrorCode::REQUEST_CANCELLED, + format!("LSP request was aborted or superseded by a newer action for URI {uri}"), + ), + + CoreError::InternalError { uri } => async_lsp::ResponseError::new( + async_lsp::ErrorCode::INTERNAL_ERROR, + format!( + "Protols internal panic: Batch processing engine lost requested file map entry {uri}" + ), + ), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub struct DocumentVersion(pub i32); + +impl DocumentVersion { + pub const DISK: Self = DocumentVersion(-1); +} + +#[derive(Clone)] +pub enum CacheEntry { + Ready { + version: DocumentVersion, + result: Result, CoreError>, + }, + Pending { + version: DocumentVersion, + future: Shared, CoreError>>>, + cancel_token: CancellationToken, + }, + Dirty, +} + +impl CacheEntry { + fn is_transient_error(result: &Result, CoreError>) -> bool { + match result { + Ok(_) + | Err( + CoreError::DiskIoError { .. } + | CoreError::InvalidPath { .. } + | CoreError::InvalidExtension { .. } + | CoreError::SourceNotFound { .. } + | CoreError::InternalError { .. }, + ) => false, + Err( + CoreError::DocumentNotFound { .. } + | CoreError::ParserError { .. } + | CoreError::RequestCancelled { .. } + | CoreError::FileTooLarge { .. }, + ) => true, + } + } + + pub fn finalize( + &mut self, + version: DocumentVersion, + parse_result: Result, CoreError>, + ) { + if let Self::Pending { + version: current_version, + .. + } = self + && version == *current_version + { + *self = Self::Ready { + version, + result: parse_result, + }; + } + } + + #[inline] + pub fn cancel(&self) { + if let Self::Pending { cancel_token, .. } = self { + cancel_token.cancel(); + } + } +} + +pub trait DocumentsGuardExt { + /// Forces the document state to `Dirty`, cancels any pending background tasks, + /// and returns `true` if the document already existed in the cache. + fn set_dirty(&mut self, uri: &Url) -> bool; + + /// Completely removes the document from the cache, triggers `cancel()` if it was `Pending`, + /// and returns `Some(CacheEntry)` if the document existed. + fn cancel_and_remove(&mut self, uri: &Url) -> Option; +} + +impl DocumentsGuardExt for RwLockWriteGuard<'_, HashMap> { + fn set_dirty(&mut self, uri: &Url) -> bool { + if let Some(entry) = self.get_mut(uri) { + entry.cancel(); + *entry = CacheEntry::Dirty; + + return true; + } + + self.insert(uri.clone(), CacheEntry::Dirty).is_some() + } + + fn cancel_and_remove(&mut self, uri: &Url) -> Option { + if let Some(entry) = self.remove(uri) { + entry.cancel(); + Some(entry) + } else { + None + } + } +} + +#[derive(Clone)] +pub struct SourceEntry { + pub version: DocumentVersion, + pub content: Arc, +} + +#[derive(Clone)] pub struct ProtoLanguageState { - sources: Arc>>, - documents: Arc>>, + client_capabilities: OnceLock, + pub(super) sources: Arc>>, + pub(super) documents: Arc>>, parser: Arc>, parsed_workspaces: Arc>>, - metamodel_query: Query, + metamodel_query: Arc, + is_indexing: Arc, +} + +const WRITE_CHUNK_SIZE: usize = 16; + +enum TaskSource { + FromMemory(Arc), + FromDisk, +} + +pub struct FinalizerToken { + inner: tokio::task::JoinHandle<()>, +} + +impl FinalizerToken { + /// Creates a dummy completed finalizer token that resolves instantly. + pub fn dummy() -> Self { + Self { + inner: tokio::spawn(async {}), + } + } + + pub(crate) fn from_handle(task_handle: tokio::task::JoinHandle<()>) -> Self { + Self { inner: task_handle } + } +} + +impl Future for FinalizerToken { + type Output = Result<(), tokio::task::JoinError>; + + #[inline] + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + Pin::new(&mut self.inner).poll(cx) + } } impl ProtoLanguageState { @@ -48,25 +274,397 @@ impl ProtoLanguageState { .expect("Tree-sitter query compilation failed"); Self { + client_capabilities: OnceLock::new(), sources: Arc::default(), documents: Arc::default(), parser: Arc::new(Mutex::new(ProtoParser::new())), parsed_workspaces: Arc::new(RwLock::new(HashSet::new())), - metamodel_query, + metamodel_query: Arc::new(metamodel_query), + is_indexing: Arc::new(AtomicBool::new(false)), } } + #[inline] + pub fn set_client_capabilities(&self, capabilities: ClientCapabilitiesSummary) { + self.client_capabilities + .set(capabilities) + .expect("Attempted to initialize client capabilities more than once per session!"); + } + + #[inline] + pub fn client_capabilities(&self) -> ClientCapabilitiesSummary { + self.client_capabilities.get().copied().unwrap_or_default() + } + + #[inline] + pub fn set_indexing(&self, value: bool) { + self.is_indexing.store(value, Ordering::Relaxed); + } + + #[inline] + pub fn is_indexing(&self) -> bool { + self.is_indexing.load(Ordering::Relaxed) + } + pub fn get_content(&self, uri: &Url) -> String { self.sources .read() .expect("poison") .get(uri) - .cloned() + .map(|e| e.content.clone()) .unwrap_or_default() } - pub fn get_document(&self, uri: &Url) -> Option { - self.documents.read().expect("poison").get(uri).cloned() + pub async fn write_lock_all( + &self, + ) -> ( + tokio::sync::RwLockWriteGuard<'_, HashMap>, + tokio::sync::RwLockWriteGuard<'_, HashMap>, + ) { + let documents_guard = self.documents.write().await; + let sources_guard = self.sources.write().await; + + (documents_guard, sources_guard) + } + + /// Processes a batch of documents. + /// + /// # Cancellation safety + /// + /// This method is cancel safe. + pub async fn query_documents_batch( + &self, + uris: UniqueUris, + batch_cancel_token: CancellationToken, + progress_reporter: Option, + need_results: bool, + ) -> ( + FinalizerToken, + impl Stream, CoreError>>, + ) { + let mut cache_hits = Vec::new(); + let mut client_futures = FuturesUnordered::new(); + let mut needs_parsing = false; + + { + let documents = self.documents.read().await; + + for uri in &uris { + match documents.get(uri) { + Some(CacheEntry::Ready { result, .. }) + if !CacheEntry::is_transient_error(result) => + { + cache_hits.extend(need_results.then(|| result.clone())); + } + Some(CacheEntry::Pending { future, .. }) => { + client_futures.extend(need_results.then(|| future.clone())); + } + _ => { + needs_parsing = true; + break; + } + } + } + } + + if !needs_parsing { + return ( + FinalizerToken::dummy(), + Either::Left(futures::stream::iter(cache_hits).chain(client_futures)), + ); + } + + cache_hits.clear(); + client_futures.clear(); + + let finalizer_futures = FuturesUnordered::new(); + let mut tasks = Vec::new(); + + { + let mut documents = self.documents.write().await; + let sources = self.sources.read().await; + + for uri in uris { + let entry = documents.get(&uri); + + match entry { + Some(CacheEntry::Ready { result, .. }) + if !CacheEntry::is_transient_error(result) => + { + cache_hits.extend(need_results.then(|| result.clone())); + continue; + } + Some(CacheEntry::Pending { future, .. }) => { + client_futures.extend(need_results.then(|| future.clone())); + continue; + } + _ => {} + } + + let source = sources.get(&uri); + + if source.is_none() + && matches!( + entry, + Some(CacheEntry::Dirty) | Some(CacheEntry::Ready { .. }) + ) + { + let result = Err(CoreError::SourceNotFound { uri: uri.clone() }); + + cache_hits.extend(need_results.then(|| result.clone())); + documents.insert( + uri.clone(), + CacheEntry::Ready { + version: DocumentVersion::DISK, + result, + }, + ); + + continue; + } + + let (expected_version, task_source) = source.map_or_else( + || (DocumentVersion::DISK, TaskSource::FromDisk), + |se| (se.version, TaskSource::FromMemory(se.content.clone())), + ); + + let (tx, rx) = oneshot::channel(); + let cancel_token = batch_cancel_token.child_token(); + let cancel_token_clone = cancel_token.clone(); + let uri_clone = uri.clone(); + + let future = async move { + let uri = uri_clone; + cancel_token_clone + .run_until_cancelled(rx) + .await + .map(|res| { + res.unwrap_or(Err(CoreError::RequestCancelled { uri: uri.clone() })) + }) + .unwrap_or(Err(CoreError::RequestCancelled { uri })) + } + .boxed() + .shared(); + + let version = expected_version; + + finalizer_futures + .push(future.clone().map(move |result| (expected_version, result))); + + client_futures.extend(need_results.then(|| future.clone())); + + documents.insert( + uri.clone(), + CacheEntry::Pending { + version, + future, + cancel_token: cancel_token.clone(), + }, + ); + + tasks.push((uri, task_source, tx, cancel_token)); + } + } + + self.run_blocking_parser(tasks); + + ( + self.run_cache_finalizer(finalizer_futures, progress_reporter), + Either::Right(futures::stream::iter(cache_hits).chain(client_futures)), + ) + } + + fn run_blocking_parser( + &self, + tasks: Vec<( + Url, + TaskSource, + oneshot::Sender, CoreError>>, + CancellationToken, + )>, + ) { + if tasks.is_empty() { + return; + } + + let query = self.metamodel_query.clone(); + tokio::task::spawn_blocking(move || { + let mut parser = tree_sitter::Parser::new(); + + let trace_error = |error: &LanguageError| { + tracing::error!( + ?error, + "Critical initialization failure: Failed to set Tree-sitter Protobuf language parser" + ); + }; + + parser + .set_language(&tree_sitter_proto::LANGUAGE.into()) + .inspect_err(trace_error) + .expect("Tree-sitter parser: setting language failed"); + + for (uri, task_source, tx, cancel_token) in tasks { + if cancel_token.is_cancelled() { + let _ = tx.send(Err(CoreError::RequestCancelled { uri })); + continue; + } + + let mut disk_buffer = Vec::new(); + + if let TaskSource::FromDisk = &task_source { + let path = match uri.to_file_path() { + Ok(path) => path, + Err(_) => { + let _ = tx.send(Err(CoreError::InvalidPath { uri })); + continue; + } + }; + + if !path.extension().is_some_and(|ext| ext == "proto") { + let _ = tx.send(Err(CoreError::InvalidExtension { uri })); + continue; + } + + match std::fs::metadata(&path) { + Err(e) => { + let _ = tx.send(Err(CoreError::DiskIoError { + uri, + details: e.to_string(), + })); + continue; + } + Ok(metadata) => { + let size = metadata.len(); + + if size > MAX_PROTO_FILE_SIZE_BYTES { + tracing::warn!( + %uri, + size, + "Refusing to parse file: size exceeds maximum limit of {} bytes", + MAX_PROTO_FILE_SIZE_BYTES + ); + let _ = tx.send(Err(CoreError::FileTooLarge { uri, size })); + continue; + } + + if !metadata.is_file() { + let _ = tx.send(Err(CoreError::InvalidPath { uri })); + continue; + } + } + } + + match std::fs::read(&path) { + Ok(bytes) => { + disk_buffer = bytes; + } + Err(error) => { + let _ = tx.send(Err(CoreError::DiskIoError { + uri, + details: error.to_string(), + })); + continue; + } + } + } + + let bytes_ref: &[u8] = match &task_source { + TaskSource::FromMemory(content) => content.as_bytes(), + TaskSource::FromDisk => &disk_buffer, + }; + + if cancel_token.is_cancelled() { + let _ = tx.send(Err(CoreError::RequestCancelled { uri })); + continue; + } + + let result = + ProtoDocument::try_from_input(uri.clone(), bytes_ref, &query, &mut parser) + .map(Arc::new) + .ok_or(CoreError::ParserError { uri }); + + let _ = tx.send(result); + } + }); + } + + fn run_cache_finalizer( + &self, + futures: FuturesUnordered, + progress_reporter: Option, + ) -> FinalizerToken + where + Fut: Future, CoreError>)> + + Send + + 'static, + { + let total_tasks = futures.len(); + + if total_tasks == 0 { + tracing::warn!( + "Document cache finalizer called with 0 tasks. This might indicate lost futures." + ); + + return FinalizerToken::dummy(); + } + + let mut completed_tasks = 0; + let mut chunks = futures.ready_chunks(WRITE_CHUNK_SIZE); + let documents_clone = self.documents.clone(); + + let inner = tokio::spawn(async move { + while let Some(chunk) = chunks.next().await { + { + let mut documents = documents_clone.write().await; + + for (expected_version, result) in chunk { + if let Some(document) = documents.get_mut(get_uri(&result)) { + document.finalize(expected_version, result); + } + + completed_tasks += 1; + } + } + + let percentage = (completed_tasks * 100 / total_tasks) as u32; + progress_reporter.report( + percentage, + &format!("Parsing [{completed_tasks}/{total_tasks}]"), + ); + + tokio::task::yield_now().await; + } + }); + + FinalizerToken { inner } + } + + pub async fn get_document( + &self, + uri: &Url, + cancel_token: CancellationToken, + ) -> Result, CoreError> { + let uris = UniqueUris::new(uri.clone()); + + let (_, mut stream) = self + .query_documents_batch(uris, cancel_token, None, true) + .await; + + match stream.next().await { + Some(result) if get_uri(&result) == uri => result, + + Some(result) => { + let returned_uri = get_uri(&result); + tracing::error!( + ?returned_uri, + expected_uri = ?uri, + "Bug: query_documents_batch returned a mismatched URL" + ); + Err(CoreError::InternalError { uri: uri.clone() }) + } + + None => Err(CoreError::DocumentNotFound { uri: uri.clone() }), + } } pub fn get_documents(&self) -> Vec { @@ -344,29 +942,6 @@ impl ProtoLanguageState { self.documents.write().expect("poison").remove(uri); } - pub fn rename_file(&mut self, new_uri: &Url, old_uri: &Url) { - info!(%new_uri, %old_uri, "renaming file"); - - let content = self.sources.write().expect("poison").remove(old_uri); - if let Some(v) = content { - self.sources - .write() - .expect("poison") - .insert(new_uri.clone(), v); - } - - let mut document = self.documents.write().expect("poison").remove(old_uri); - if let Some(ref mut v) = document { - v.uri = new_uri.clone(); - } - if let Some(v) = document { - self.documents - .write() - .expect("poison") - .insert(new_uri.clone(), v); - } - } - pub fn completion_items_for_document(&self, url: &Url) -> Vec { let collector = |f: fn(&ElementKind) -> bool, k: CompletionItemKind| { self.get_document(url) @@ -429,6 +1004,23 @@ fn is_message_kind(kind: &ElementKind) -> bool { matches!(kind, ElementKind::Message { .. }) } +fn get_uri(result: &Result, CoreError>) -> &Url { + match result { + Ok(d) => &d.uri, + Err( + CoreError::DiskIoError { uri, .. } + | CoreError::DocumentNotFound { uri } + | CoreError::InternalError { uri } + | CoreError::InvalidPath { uri } + | CoreError::InvalidExtension { uri } + | CoreError::FileTooLarge { uri, .. } + | CoreError::ParserError { uri } + | CoreError::RequestCancelled { uri } + | CoreError::SourceNotFound { uri }, + ) => uri, + } +} + #[cfg(test)] mod test { use super::*; @@ -606,24 +1198,24 @@ mod test { assert_eq!(state.get_content(&test_uri), ""); } - #[test] - fn test_rename_file() { - let mut state = setup_state(); - let old_uri = uri("file:///test.proto"); - let new_uri = uri("file:///renamed.proto"); + // #[test] + // fn test_rename_file() { + // let mut state = setup_state(); + // let old_uri = uri("file:///test.proto"); + // let new_uri = uri("file:///renamed.proto"); - assert!(state.get_document(&old_uri).is_some()); - assert!(state.get_document(&new_uri).is_none()); + // assert!(state.get_document(&old_uri).is_some()); + // assert!(state.get_document(&new_uri).is_none()); - state.rename_file(&new_uri, &old_uri); + // state.rename_file(&new_uri, &old_uri); - assert!(state.get_document(&old_uri).is_none()); - assert!(state.get_document(&new_uri).is_some()); - assert_eq!( - state.get_content(&new_uri), - "syntax = \"proto3\";\npackage com.test;\nmessage Book { string title = 1; }\nenum Color { RED = 0; }\n" - ); - } + // assert!(state.get_document(&old_uri).is_none()); + // assert!(state.get_document(&new_uri).is_some()); + // assert_eq!( + // state.get_content(&new_uri), + // "syntax = \"proto3\";\npackage com.test;\nmessage Book { string title = 1; }\nenum Color { RED = 0; }\n" + // ); + // } #[test] fn test_upsert_content_tracks_unresolved_imports() { diff --git a/src/state/rename.rs b/src/state/rename.rs index 549dbc6..4969fbe 100644 --- a/src/state/rename.rs +++ b/src/state/rename.rs @@ -2,6 +2,7 @@ use std::collections::{BTreeMap, HashMap}; use std::path::PathBuf; use async_lsp::lsp_types::{Location, Position, TextEdit, Url}; +use tokio_util::sync::CancellationToken; use crate::model::ElementKind; use crate::state::ProtoLanguageState; @@ -69,13 +70,14 @@ impl ProtoLanguageState { /// `chain_rpc_request_response` gates the rpc/request/response chain: when /// `false`, only the primary op is returned. It is wired to the /// `[config.rename]` `chain_rpc_request_response` setting. - pub fn compute_rename_ops( + pub async fn compute_rename_ops( &self, decl_uri: &Url, decl_pos: Position, new_name: &str, ipath: &[PathBuf], chain_rpc_request_response: bool, + cancel_token: CancellationToken, ) -> Vec { let mut ops = vec![RenameOp { uri: decl_uri.clone(), @@ -83,7 +85,10 @@ impl ProtoLanguageState { new_name: new_name.to_owned(), }]; if chain_rpc_request_response { - ops.extend(self.compute_chain_siblings(decl_uri, decl_pos, new_name, ipath)); + ops.extend( + self.compute_chain_siblings(decl_uri, decl_pos, new_name, ipath, cancel_token) + .await, + ); } ops } @@ -92,10 +97,14 @@ impl ProtoLanguageState { /// single map. Returns `None` if the *primary* (first) op fails — in that /// case the user's invocation should produce no edit at all. Sibling /// failures are silently skipped so the primary always lands. - pub fn apply_rename_ops(&mut self, ops: &[RenameOp]) -> Option>> { + pub async fn apply_rename_ops( + &mut self, + ops: &[RenameOp], + cancel_token: CancellationToken, + ) -> Option>> { let mut all: HashMap> = HashMap::new(); for (i, op) in ops.iter().enumerate() { - match self.run_single_rename(op) { + match self.run_single_rename(op, cancel_token.clone()).await { Some(edits) => { for (u, e) in edits { all.entry(u).or_default().extend(e); @@ -108,49 +117,64 @@ impl ProtoLanguageState { Some(all) } - fn run_single_rename(&mut self, op: &RenameOp) -> Option>> { + async fn run_single_rename( + &mut self, + op: &RenameOp, + cancel_token: CancellationToken, + ) -> Option>> { // The workspace is already fully indexed once at startup (see the LSP // `initialize` handler), so cross-file rename resolves against the // cached metamodel pool without any per-request re-scan. - let target_fqn = self.resolve_target_fqn(&op.uri, op.pos)?; + let target_fqn = self + .resolve_target_fqn(&op.uri, op.pos, cancel_token) + .await?; Some(self.rename_for_fqn(&target_fqn, &op.new_name)) } - fn compute_chain_siblings( + async fn compute_chain_siblings( &self, decl_uri: &Url, decl_pos: Position, new_name: &str, ipath: &[PathBuf], + cancel_token: CancellationToken, ) -> Vec { if new_name.is_empty() { return vec![]; } - let Some(document) = self.get_document(decl_uri) else { + let Ok(document) = self.get_document(decl_uri, cancel_token.clone()).await else { return vec![]; }; let content = self.get_content(decl_uri); let bytes = content.as_bytes(); if document.rpc_at_position(decl_pos, bytes).is_some() { - return self.chain_from_rpc_cursor(decl_uri, decl_pos, new_name, ipath); + return self + .chain_from_rpc_cursor(decl_uri, decl_pos, new_name, ipath, cancel_token) + .await; } if document.message_name_at_position(decl_pos, bytes).is_some() { - return self.chain_from_message_cursor(decl_uri, decl_pos, new_name, ipath); + return self + .chain_from_message_cursor(decl_uri, decl_pos, new_name, ipath, cancel_token) + .await; } vec![] } /// Case A: user invoked rename on an `rpc_name`. The primary is the rpc /// rename; siblings are the convention-matching request/response messages. - fn chain_from_rpc_cursor( + async fn chain_from_rpc_cursor( &self, decl_uri: &Url, decl_pos: Position, new_name: &str, ipath: &[PathBuf], + cancel_token: CancellationToken, ) -> Vec { - let document = self.get_document(decl_uri).expect("checked by caller"); + let document = self + .get_document(decl_uri, cancel_token.clone()) + .await + .expect("checked by caller"); let content = self.get_content(decl_uri); let (old_rpc_name, request_text, response_text) = document .rpc_at_position(decl_pos, content.as_bytes()) @@ -164,7 +188,9 @@ impl ProtoLanguageState { ("Request", request_text.as_str()), ("Response", response_text.as_str()), ], + cancel_token, ) + .await } /// Case B: user invoked rename on a `message_name` matching the @@ -176,14 +202,19 @@ impl ProtoLanguageState { /// unique one whose request/response slot resolves (via the workspace's /// name-resolution rules) to the user's primary message. If zero or more /// than one rpc matches, the chain is silently dropped. - fn chain_from_message_cursor( + async fn chain_from_message_cursor( &self, decl_uri: &Url, decl_pos: Position, new_name: &str, ipath: &[PathBuf], + cancel_token: CancellationToken, ) -> Vec { - let document = self.get_document(decl_uri).expect("checked by caller"); + let cancel_token_clone = cancel_token.clone(); + let document = self + .get_document(decl_uri, cancel_token_clone) + .await + .expect("checked by caller"); let content = self.get_content(decl_uri); let msg_name = document .message_name_at_position(decl_pos, content.as_bytes()) @@ -202,7 +233,8 @@ impl ProtoLanguageState { // slot resolves to the user's actual declaration. let mut matching: Vec<(Location, String, String)> = vec![]; for rpc_loc in self.find_rpc_decls(&rpc_base) { - let Some(rpc_document) = self.get_document(&rpc_loc.uri) else { + let Ok(rpc_document) = self.get_document(&rpc_loc.uri, cancel_token.clone()).await + else { continue; }; let rpc_content = self.get_content(&rpc_loc.uri); @@ -248,27 +280,32 @@ impl ProtoLanguageState { } else { ("Request", rpc_req.as_str()) }; - ops.extend(self.sibling_message_ops( - &rpc_loc.uri, - &rpc_base, - &new_rpc_base, - ipath, - &[(opposite_suffix, opposite_text)], - )); + ops.extend( + self.sibling_message_ops( + &rpc_loc.uri, + &rpc_base, + &new_rpc_base, + ipath, + &[(opposite_suffix, opposite_text)], + cancel_token, + ) + .await, + ); ops } /// Resolve each `(suffix, type_text)` slot to a message declaration and /// build the corresponding rename op, gated on convention + uniqueness. - fn sibling_message_ops( + async fn sibling_message_ops( &self, anchor_uri: &Url, old_rpc_name: &str, new_rpc_name: &str, _ipath: &[PathBuf], slots: &[(&str, &str)], + cancel_token: CancellationToken, ) -> Vec { - let Some(anchor_document) = self.get_document(anchor_uri) else { + let Ok(anchor_document) = self.get_document(anchor_uri, cancel_token).await else { return vec![]; }; let anchor_package = anchor_document.package_name().to_owned(); @@ -319,6 +356,7 @@ mod test { use async_lsp::lsp_types::{Position, Url}; use insta::assert_yaml_snapshot; + use tokio_util::sync::CancellationToken; use crate::config::Config; use crate::state::ProtoLanguageState; @@ -384,8 +422,8 @@ mod test { ); } - #[test] - fn test_resolve_target_fqn() { + #[tokio::test(flavor = "current_thread")] + async fn test_resolve_target_fqn() { let ipath = vec![PathBuf::from("src/state/input")]; let state = make_state( &[ @@ -399,35 +437,44 @@ mod test { // Cursor on the `Author` message declaration name in b.proto. assert_eq!( - state.resolve_target_fqn( - &b_uri, - Position { - line: 5, - character: 10 - } - ), + state + .resolve_target_fqn( + &b_uri, + Position { + line: 5, + character: 10 + }, + CancellationToken::new() + ) + .await, Some("com.workspace.Author".to_owned()) ); // Cursor on the `Author` type reference inside a field in a.proto. assert_eq!( - state.resolve_target_fqn( - &a_uri, - Position { - line: 11, - character: 5 - } - ), + state + .resolve_target_fqn( + &a_uri, + Position { + line: 11, + character: 5 + }, + CancellationToken::new() + ) + .await, Some("com.workspace.Author".to_owned()) ); // Cursor on whitespace -> None. assert_eq!( - state.resolve_target_fqn( - &a_uri, - Position { - line: 0, - character: 0 - } - ), + state + .resolve_target_fqn( + &a_uri, + Position { + line: 0, + character: 0 + }, + CancellationToken::new() + ) + .await, None ); } @@ -465,8 +512,8 @@ mod test { assert_eq!(state.count_rpc_uses_of_type("Unrelated"), 0); } - #[test] - fn test_compute_rename_ops_cross_package_collision_blocks_chain() { + #[tokio::test(flavor = "current_thread")] + async fn test_compute_rename_ops_cross_package_collision_blocks_chain() { // Two services in different packages each declare `rpc GetBook`. Only // pkg foo follows the convention; pkg bar uses unrelated message // names. Without the per-candidate full-qualified resolution check, @@ -494,7 +541,16 @@ mod test { line: 4, character: 12, }; - let ops = state.compute_rename_ops(&foo_uri, pos, "FetchBookRequest", &ipath, true); + let ops = state + .compute_rename_ops( + &foo_uri, + pos, + "FetchBookRequest", + &ipath, + true, + CancellationToken::new(), + ) + .await; // Primary + foo's rpc + foo's response. bar's GetBook is NOT touched. assert_eq!(ops.len(), 3, "expected exactly the foo trio, got {ops:?}"); @@ -506,8 +562,8 @@ mod test { } } - #[test] - fn test_compute_rename_ops_chain_from_rpc_cursor() { + #[tokio::test(flavor = "current_thread")] + async fn test_compute_rename_ops_chain_from_rpc_cursor() { let ipath = vec![PathBuf::from("src/state/input")]; let svc_uri = "file://input/service.proto".parse().unwrap(); let state = make_state( @@ -529,7 +585,16 @@ mod test { line: 7, character: 10, }; - let ops = state.compute_rename_ops(&svc_uri, pos, "FetchBook", &ipath, true); + let ops = state + .compute_rename_ops( + &svc_uri, + pos, + "FetchBook", + &ipath, + true, + CancellationToken::new(), + ) + .await; // Primary + Request + Response. assert_eq!( @@ -548,8 +613,8 @@ mod test { ); } - #[test] - fn test_compute_rename_ops_chain_disabled() { + #[tokio::test(flavor = "current_thread")] + async fn test_compute_rename_ops_chain_disabled() { // Same setup as `chain_from_rpc_cursor`, but with the chain flag off: // the rpc/request/response chain is gated behind the `[config.rename]` // `chain_rpc_request_response` setting, so only the primary op fires. @@ -574,7 +639,16 @@ mod test { line: 7, character: 10, }; - let ops = state.compute_rename_ops(&svc_uri, pos, "FetchBook", &ipath, false); + let ops = state + .compute_rename_ops( + &svc_uri, + pos, + "FetchBook", + &ipath, + false, + CancellationToken::new(), + ) + .await; assert_eq!( ops.len(), @@ -584,8 +658,8 @@ mod test { assert_eq!(ops[0], op("file://input/service.proto", 7, 10, "FetchBook")); } - #[test] - fn test_compute_rename_ops_chain_from_request_cursor() { + #[tokio::test(flavor = "current_thread")] + async fn test_compute_rename_ops_chain_from_request_cursor() { let ipath = vec![PathBuf::from("src/state/input")]; let msg_uri = "file://input/messages.proto".parse().unwrap(); let state = make_state( @@ -607,7 +681,16 @@ mod test { line: 4, character: 12, }; - let ops = state.compute_rename_ops(&msg_uri, pos, "FetchBookRequest", &ipath, true); + let ops = state + .compute_rename_ops( + &msg_uri, + pos, + "FetchBookRequest", + &ipath, + true, + CancellationToken::new(), + ) + .await; // Primary (request) + rpc + response. assert_eq!( @@ -626,8 +709,8 @@ mod test { ); } - #[test] - fn test_compute_rename_ops_shared_request_blocks_chain() { + #[tokio::test(flavor = "current_thread")] + async fn test_compute_rename_ops_shared_request_blocks_chain() { let ipath = vec![PathBuf::from("src/state/input")]; let msg_uri = "file://input/messages.proto".parse().unwrap(); let state = make_state( @@ -650,7 +733,16 @@ mod test { line: 10, character: 12, }; - let ops = state.compute_rename_ops(&msg_uri, pos, "RenamedReq", &ipath, true); + let ops = state + .compute_rename_ops( + &msg_uri, + pos, + "RenamedReq", + &ipath, + true, + CancellationToken::new(), + ) + .await; assert_eq!( ops.len(), 1, @@ -662,8 +754,8 @@ mod test { ); } - #[test] - fn test_compute_rename_ops_new_name_breaks_convention() { + #[tokio::test(flavor = "current_thread")] + async fn test_compute_rename_ops_new_name_breaks_convention() { let ipath = vec![PathBuf::from("src/state/input")]; let msg_uri = "file://input/messages.proto".parse().unwrap(); let state = make_state( @@ -686,7 +778,16 @@ mod test { line: 4, character: 12, }; - let ops = state.compute_rename_ops(&msg_uri, pos, "Whatever", &ipath, true); + let ops = state + .compute_rename_ops( + &msg_uri, + pos, + "Whatever", + &ipath, + true, + CancellationToken::new(), + ) + .await; assert_eq!( ops.len(), 1, @@ -695,8 +796,8 @@ mod test { assert_eq!(ops[0], op("file://input/messages.proto", 4, 12, "Whatever")); } - #[test] - fn test_compute_rename_ops_reference_site_pivot_unsupported() { + #[tokio::test(flavor = "current_thread")] + async fn test_compute_rename_ops_reference_site_pivot_unsupported() { // compute_rename_ops requires the caller to have pivoted to the // declaration. As a sanity check, calling it with a cursor on a // reference site (not a declaration) yields a primary-only op that @@ -724,14 +825,23 @@ mod test { line: 7, character: 20, }; - let ops = state.compute_rename_ops(&svc_uri, pos, "RenamedRequest", &ipath, true); + let ops = state + .compute_rename_ops( + &svc_uri, + pos, + "RenamedRequest", + &ipath, + true, + CancellationToken::new(), + ) + .await; // Single op (the primary at the reference site) — no chain. This // documents the contract: the LSP layer must pivot first. assert_eq!(ops.len(), 1, "{ops:?}"); } - #[test] - fn test_apply_rename_ops_chain_from_rpc_cursor() { + #[tokio::test(flavor = "current_thread")] + async fn test_apply_rename_ops_chain_from_rpc_cursor() { // End-to-end snapshot of the full `rpc (Request) returns // (Response)` convention chain. Renaming `GetBook` → // `FetchBook` should fan out into: @@ -763,9 +873,19 @@ mod test { line: 7, character: 10, }; - let ops = state.compute_rename_ops(&svc_uri, pos, "FetchBook", &ipath, true); + let ops = state + .compute_rename_ops( + &svc_uri, + pos, + "FetchBook", + &ipath, + true, + CancellationToken::new(), + ) + .await; let edits = state - .apply_rename_ops(&ops) + .apply_rename_ops(&ops, CancellationToken::new()) + .await .expect("primary rename should not fail"); // Sort within each file so the snapshot is order-independent across @@ -907,8 +1027,8 @@ mod test { ); } - #[test] - fn test_apply_rename_ops_from_reference_site() { + #[tokio::test(flavor = "current_thread")] + async fn test_apply_rename_ops_from_reference_site() { // End-to-end: invoking rename on a *reference site* (the `GetBookRequest` // type inside the rpc signature) pivots to the declaration and renames // both the declaration and the reference. @@ -934,8 +1054,20 @@ mod test { line: 7, character: 19, }; - let ops = state.compute_rename_ops(&svc_uri, pos, "FetchBookRequest", &ipath, false); - let edits = state.apply_rename_ops(&ops).expect("rename should succeed"); + let ops = state + .compute_rename_ops( + &svc_uri, + pos, + "FetchBookRequest", + &ipath, + false, + CancellationToken::new(), + ) + .await; + let edits = state + .apply_rename_ops(&ops, CancellationToken::new()) + .await + .expect("rename should succeed"); let mut normalized: std::collections::BTreeMap> = std::collections::BTreeMap::new(); for (url, mut v) in edits { diff --git a/src/state/resolve.rs b/src/state/resolve.rs index b8bcad0..47ccd36 100644 --- a/src/state/resolve.rs +++ b/src/state/resolve.rs @@ -8,6 +8,7 @@ use std::collections::BTreeMap; use async_lsp::lsp_types::{Location, Position, TextEdit, Url}; +use tokio_util::sync::CancellationToken; use crate::model::{ModelElement, SpatialEntry, TypeReference}; use crate::state::ProtoLanguageState; @@ -93,8 +94,13 @@ impl ProtoLanguageState { /// returned. If it rests on a type reference, the segment under the cursor /// is resolved to its referenced definition's FQN (so renaming the outer /// segment of `Book.Author` targets `Book`, not the nested `Author`). - pub fn resolve_target_fqn(&self, uri: &Url, position: Position) -> Option { - let document = self.get_document(uri)?; + pub async fn resolve_target_fqn( + &self, + uri: &Url, + position: Position, + cancel_token: CancellationToken, + ) -> Option { + let document = self.get_document(uri, cancel_token).await.ok()?; let SpatialEntry { element_id, .. } = document.find_entry_at_position(position)?; let element = document.elements.get(*element_id)?; diff --git a/src/state/unique_uris.rs b/src/state/unique_uris.rs new file mode 100644 index 0000000..53b5ca9 --- /dev/null +++ b/src/state/unique_uris.rs @@ -0,0 +1,95 @@ +use std::{collections::HashSet, ops::Deref}; + +use async_lsp::lsp_types::Url; + +/// A type-safe container that always holds unique and sorted URLs. +pub struct UniqueUris { + uris: Vec, +} + +impl UniqueUris { + #[inline] + pub fn new(source: T) -> Self + where + T: IntoUniqueUris, + { + source.into_unique_uris() + } + + #[inline] + pub fn iter(&self) -> std::slice::Iter<'_, Url> { + self.uris.iter() + } + + #[inline] + pub fn into_inner(self) -> Vec { + self.uris + } +} + +impl Deref for UniqueUris { + type Target = [Url]; + + #[inline] + fn deref(&self) -> &Self::Target { + &self.uris + } +} + +impl IntoIterator for UniqueUris { + type Item = Url; + type IntoIter = std::vec::IntoIter; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.uris.into_iter() + } +} + +impl<'a> IntoIterator for &'a UniqueUris { + type Item = &'a Url; + type IntoIter = std::slice::Iter<'a, Url>; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.uris.iter() + } +} + +pub trait IntoUniqueUris { + fn into_unique_uris(self) -> UniqueUris; +} + +impl IntoUniqueUris for HashSet { + #[inline] + fn into_unique_uris(self) -> UniqueUris { + let uris: Vec = self.into_iter().collect(); + UniqueUris { uris } + } +} + +impl IntoUniqueUris for std::collections::hash_map::IntoKeys { + #[inline] + fn into_unique_uris(self) -> UniqueUris { + UniqueUris { + uris: self.collect(), + } + } +} + +impl IntoUniqueUris for Vec { + #[inline] + fn into_unique_uris(self) -> UniqueUris { + let mut uris = self; + uris.sort_unstable(); + uris.dedup(); + UniqueUris { uris } + } +} + +impl IntoUniqueUris for Url { + #[inline] + fn into_unique_uris(self) -> UniqueUris { + UniqueUris { uris: vec![self] } + } +} diff --git a/src/transport.rs b/src/transport.rs index 8bdbe00..bf5cfa9 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -138,7 +138,7 @@ async fn create_pipe_transport(path: &str) -> TransportResult<(LspReader, LspWri #[cfg(any(windows, test))] fn normalize_windows_pipe(path: &str) -> TransportResult { let full_path = match path { - local if local.starts_with(r"\\.\pipe\") => local.to_string(), + local if local.starts_with(r"\\.\pipe\") => local.to_owned(), unc if unc.starts_with(r"\\") => { let mut components = unc.split('\\').skip(2); @@ -146,7 +146,7 @@ fn normalize_windows_pipe(path: &str) -> TransportResult { let pipe_segment = components.next(); if pipe_segment == Some("pipe") { - unc.to_string() + unc.to_owned() } else { return Err(format!( "Invalid UNC pipe path: '{unc}'. Missing or misplaced '\\pipe\\'." @@ -201,7 +201,7 @@ async fn test_unix_socket_wont_delete_regular_file() { let path = file.path().to_str().unwrap(); let cli = Cli { - pipe: Some(path.to_string()), + pipe: Some(path.to_owned()), ..Default::default() }; diff --git a/src/utils.rs b/src/utils.rs index 4c9c70c..ce33635 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,4 +1,7 @@ +use std::path::PathBuf; + use async_lsp::lsp_types::{Position, Range}; +use futures::FutureExt; use tree_sitter::{Node, Point}; /// Converts a Tree-sitter [`Point`] into an LSP [`Position`]. @@ -108,17 +111,100 @@ pub fn clean_proto_comment(raw_text: &str) -> String { if let Some(inner) = raw_text.strip_prefix("/*") && let Some(uncommented) = inner.strip_suffix("*/") { - return uncommented.to_string(); + return uncommented.to_owned(); } if let Some(uncommented) = raw_text.strip_prefix("//") { return uncommented .strip_prefix(' ') .unwrap_or(uncommented) - .to_string(); + .to_owned(); + } + + raw_text.to_owned() +} + +pub(crate) async fn wait_for_shutdown_signals() { + let ctrl_c = async { + let _ = tokio::signal::ctrl_c().await; + tracing::info!("Ctrl+C received"); + } + .boxed(); + + let os_signal = async { + #[cfg(unix)] + { + if let Ok(mut stream) = + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + { + stream.recv().await; + tracing::info!("SIGTERM received"); + } + } + + #[cfg(windows)] + { + if let Ok(mut stream) = tokio::signal::windows::ctrl_close() { + stream.recv().await; + tracing::info!("Windows console close event received"); + } + } + + #[cfg(not(any(unix, windows)))] + futures::future::pending::<()>().await; } + .boxed(); + + let _ = futures::future::select(ctrl_c, os_signal).await; +} + +#[derive(Debug, Clone, Default)] +pub struct OutermostPaths(Vec); + +impl OutermostPaths { + pub fn from_iter(paths: impl IntoIterator) -> Self { + let mut raw_paths: Vec<_> = paths.into_iter().collect(); + + raw_paths.sort_unstable(); + raw_paths.dedup(); + + let mut outermost = Vec::with_capacity(raw_paths.len()); + let mut last_added = None; - raw_text.to_string() + for path in &raw_paths { + if !last_added.is_some_and(|outer| path.starts_with(outer)) { + outermost.push(path.clone()); + last_added = Some(path); + } + } + + Self(outermost) + } + + pub fn into_inner(self) -> Vec { + self.0 + } + + pub fn iter(&self) -> std::slice::Iter<'_, PathBuf> { + self.0.iter() + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + pub fn len(&self) -> usize { + self.0.len() + } +} + +impl IntoIterator for OutermostPaths { + type Item = PathBuf; + type IntoIter = std::vec::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } } #[cfg(test)]