Add support for workspace symbols

Cargo.toml

@@ -32,4 +32,4 @@ pkg-config = "0.3"
 clap = { version = "4.5", features = ["derive"] }
 
 [dev-dependencies]
-insta = { version = "1.43", features = ["yaml"] }
+insta = { version = "1.43", features = ["yaml", "redactions"] }

src/config/workspace.rs

@@ -131,6 +131,10 @@ impl WorkspaceProtoConfigs {
         Some(ipath)
     }
 
+    pub fn get_workspaces(&self) -> Vec<&Url> {
+        self.workspaces.iter().collect()
+    }
+
     pub fn no_workspace_mode(&mut self) {
         let wr = ProtolsConfig::default();
         let rp = if cfg!(target_os = "windows") {

src/lsp.rs

@@ -14,7 +14,7 @@ use async_lsp::lsp_types::{
     RenameParams, ServerCapabilities, ServerInfo, TextDocumentPositionParams,
     TextDocumentSyncCapability, TextDocumentSyncKind, TextEdit, Url, WorkspaceEdit,
     WorkspaceFileOperationsServerCapabilities, WorkspaceFoldersServerCapabilities,
-    WorkspaceServerCapabilities,
+    WorkspaceServerCapabilities, WorkspaceSymbolParams, WorkspaceSymbolResponse,
 };
 use async_lsp::{LanguageClient, ResponseError};
 use futures::future::BoxFuture;
@@ -113,6 +113,7 @@ impl ProtoLanguageServer {
                 definition_provider: Some(OneOf::Left(true)),
                 hover_provider: Some(HoverProviderCapability::Simple(true)),
                 document_symbol_provider: Some(OneOf::Left(true)),
+                workspace_symbol_provider: Some(OneOf::Left(true)),
                 completion_provider: Some(CompletionOptions::default()),
                 rename_provider: Some(rename_provider),
                 document_formatting_provider: Some(OneOf::Left(true)),
@@ -379,6 +380,35 @@ impl ProtoLanguageServer {
         Box::pin(async move { Ok(Some(response)) })
     }
 
+    pub(super) fn workspace_symbol(
+        &mut self,
+        params: WorkspaceSymbolParams,
+    ) -> BoxFuture<'static, Result<Option<WorkspaceSymbolResponse>, ResponseError>> {
+        let query = params.query.to_lowercase();
+        let work_done_token = params.work_done_progress_params.work_done_token;
+
+        // Parse all files from all workspaces
+        let workspaces = self.configs.get_workspaces();
+        let progress_sender = work_done_token.map(|token| self.with_report_progress(token));
+
+        for workspace in workspaces {
+            if let Ok(workspace_path) = workspace.to_file_path() {
+                self.state
+                    .parse_all_from_workspace(workspace_path, progress_sender.clone());
+            }
+        }
+
+        let symbols = self.state.find_workspace_symbols(&query);
+
+        Box::pin(async move {
+            if symbols.is_empty() {
+                Ok(None)
+            } else {
+                Ok(Some(WorkspaceSymbolResponse::Nested(symbols)))
+            }
+        })
+    }
+
     pub(super) fn formatting(
         &mut self,
         params: DocumentFormattingParams,

src/server.rs

@@ -9,6 +9,7 @@ use async_lsp::{
         request::{
             Completion, DocumentSymbolRequest, Formatting, GotoDefinition, HoverRequest,
             Initialize, PrepareRenameRequest, RangeFormatting, References, Rename,
+            WorkspaceSymbolRequest,
         },
     },
     router::Router,
@@ -59,6 +60,7 @@ impl ProtoLanguageServer {
         router.request::<References, _>(|st, params| st.references(params));
         router.request::<GotoDefinition, _>(|st, params| st.definition(params));
         router.request::<DocumentSymbolRequest, _>(|st, params| st.document_symbol(params));
+        router.request::<WorkspaceSymbolRequest, _>(|st, params| st.workspace_symbol(params));
         router.request::<Formatting, _>(|st, params| st.formatting(params));
         router.request::<RangeFormatting, _>(|st, params| st.range_formatting(params));
 

src/state.rs

@@ -5,8 +5,11 @@ use std::{
 };
 use tracing::info;
 
-use async_lsp::lsp_types::ProgressParamsValue;
-use async_lsp::lsp_types::{CompletionItem, CompletionItemKind, PublishDiagnosticsParams, Url};
+use async_lsp::lsp_types::{
+    CompletionItem, CompletionItemKind, Location, OneOf, PublishDiagnosticsParams, Url,
+    WorkspaceSymbol,
+};
+use async_lsp::lsp_types::{DocumentSymbol, ProgressParamsValue};
 use std::sync::mpsc::Sender;
 use tree_sitter::Node;
 use walkdir::WalkDir;
@@ -73,6 +76,78 @@ impl ProtoLanguageState {
             .collect()
     }
 
+    pub fn find_workspace_symbols(&self, query: &str) -> Vec<WorkspaceSymbol> {
+        let mut symbols = Vec::new();
+
+        for tree in self.get_trees() {
+            let content = self.get_content(&tree.uri);
+            let doc_symbols = tree.find_document_locations(content.as_bytes());
+
+            for doc_symbol in doc_symbols {
+                Self::find_workspace_symbols_impl(
+                    &doc_symbol,
+                    &tree.uri,
+                    query,
+                    None,
+                    &mut symbols,
+                );
+            }
+        }
+
+        // Sort symbols by name and then by URI for consistent ordering
+        symbols.sort_by(|a, b| {
+            let name_cmp = a.name.cmp(&b.name);
+            if name_cmp != std::cmp::Ordering::Equal {
+                return name_cmp;
+            }
+            // Extract URI from location
+            match (&a.location, &b.location) {
+                (OneOf::Left(loc_a), OneOf::Left(loc_b)) => {
+                    loc_a.uri.as_str().cmp(loc_b.uri.as_str())
+                }
+                _ => std::cmp::Ordering::Equal,
+            }
+        });
+
+        symbols
+    }
+
+    fn find_workspace_symbols_impl(
+        doc_symbol: &DocumentSymbol,
+        uri: &Url,
+        query: &str,
+        container_name: Option<String>,
+        symbols: &mut Vec<WorkspaceSymbol>,
+    ) {
+        let symbol_name_lower = doc_symbol.name.to_lowercase();
+
+        if query.is_empty() || symbol_name_lower.contains(query) {
+            symbols.push(WorkspaceSymbol {
+                name: doc_symbol.name.clone(),
+                kind: doc_symbol.kind,
+                tags: doc_symbol.tags.clone(),
+                container_name: container_name.clone(),
+                location: OneOf::Left(Location {
+                    uri: uri.clone(),
+                    range: doc_symbol.range,
+                }),
+                data: None,
+            });
+        }
+
+        if let Some(children) = &doc_symbol.children {
+            for child in children {
+                Self::find_workspace_symbols_impl(
+                    child,
+                    uri,
+                    query,
+                    Some(doc_symbol.name.clone()),
+                    symbols,
+                );
+            }
+        }
+    }
+
     fn upsert_content_impl(
         &mut self,
         uri: &Url,

src/workspace/definition.rs

@@ -107,13 +107,9 @@ mod test {
             Jumpable::Import("c.proto".to_owned()),
         );
 
-        assert_eq!(loc.len(), 1);
-        assert!(
-            loc[0]
-                .uri
-                .to_file_path()
-                .unwrap()
-                .ends_with(ipath[0].join("c.proto"))
-        )
+        assert_yaml_snapshot!(loc, {"[0].uri" => insta::dynamic_redaction(|c, _| {
+            assert!(c.as_str().unwrap().ends_with("c.proto"));
+            "file://<redacted>/c.proto".to_string()
+        })});
     }
 }

src/workspace/mod.rs

@@ -1,3 +1,4 @@
 mod definition;
 mod hover;
 mod rename;
+mod workspace_symbol;