refactor: reuse code and fix a didClose missing crash (#10)

Some housekeeping!

Author: asharkhan3101

Cargo.lock

@@ -1,15 +1,15 @@
 [package]
 name = "protols"
 description = "Language server for proto3 files"
-version = "0.1.1"
+version = "0.2.0"
 edition = "2021"
 license = "MIT"
 homepage = "https://github.com/coder3101/protols"
 repository = "https://github.com/coder3101/protols"
 readme = "README.md"
 keywords = ["lsp", "proto3"]
 
-exclude = ["assets/*"]
+exclude = ["assets/*", "sample/*"]
 
 [dependencies]
 async-lsp = { version = "0.2.0", features = ["tokio"] }

Cargo.toml

@@ -7,7 +7,7 @@ A Language Server for **proto3** files. It uses tree-sitter parser for all opera
 - [x] Hover
 - [x] Go to definition
 - [x] Diagnostics
-- [x] Document symbols outline for message and enums
+- [x] Document Symbols for message and enums
 
 ## Installation
 

README.md

@@ -54,3 +54,4 @@ enum EnumSample {
     STARTED = 1;
     RUNNING = 1;
 }
+

src/simple.proto renamed to sample/simple.proto

@@ -1,13 +1,14 @@
 use std::ops::ControlFlow;
-use tracing::{debug, info};
+use tracing::{error, info};
 
 use async_lsp::lsp_types::{
-    DidChangeTextDocumentParams, DidOpenTextDocumentParams, DidSaveTextDocumentParams,
-    DocumentSymbolParams, DocumentSymbolResponse, GotoDefinitionParams, GotoDefinitionResponse,
-    Hover, HoverContents, HoverParams, HoverProviderCapability, InitializeParams, InitializeResult,
-    OneOf, ServerCapabilities, ServerInfo, TextDocumentSyncCapability, TextDocumentSyncKind,
+    DidChangeTextDocumentParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams,
+    DidSaveTextDocumentParams, DocumentSymbolParams, DocumentSymbolResponse, GotoDefinitionParams,
+    GotoDefinitionResponse, Hover, HoverContents, HoverParams, HoverProviderCapability,
+    InitializeParams, InitializeResult, OneOf, ServerCapabilities, ServerInfo,
+    TextDocumentSyncCapability, TextDocumentSyncKind,
 };
-use async_lsp::{ErrorCode, LanguageClient, LanguageServer, ResponseError};
+use async_lsp::{LanguageClient, LanguageServer, ResponseError};
 use futures::future::BoxFuture;
 
 use crate::server::ServerState;
@@ -29,7 +30,6 @@ impl LanguageServer for ServerState {
         let cversion = version.unwrap_or("<unknown>");
 
         info!("Connected with client {cname} {cversion}");
-        debug!("Initialize with {params:?}");
 
         let response = InitializeResult {
             capabilities: ServerCapabilities {
@@ -58,39 +58,26 @@ impl LanguageServer for ServerState {
         let uri = param.text_document_position_params.text_document.uri;
         let pos = param.text_document_position_params.position;
 
-        let Some(contents) = self.documents.get(&uri) else {
-            return Box::pin(async move {
-                Err(ResponseError::new(
-                    ErrorCode::INVALID_REQUEST,
-                    "uri was never opened",
-                ))
-            });
-        };
-
-        let Some(parsed) = self.parser.parse(contents.as_bytes()) else {
-            return Box::pin(async move {
-                Err(ResponseError::new(
-                    ErrorCode::REQUEST_FAILED,
-                    "ts failed to parse contents",
-                ))
-            });
-        };
-
-        let comments = parsed.hover(&pos, contents.as_bytes());
-        info!("Found {} node comments in the document", comments.len());
-        let response = match comments.len() {
-            0 => None,
-            1 => Some(Hover {
-                contents: HoverContents::Scalar(comments[0].clone()),
-                range: None,
-            }),
-            2.. => Some(Hover {
-                contents: HoverContents::Array(comments),
-                range: None,
-            }),
-        };
-
-        Box::pin(async move { Ok(response) })
+        match self.get_parsed_tree_and_content(&uri) {
+            Err(e) => Box::pin(async move { Err(e) }),
+            Ok((tree, content)) => {
+                let comments = tree.hover(&pos, content.as_bytes());
+
+                let response = match comments.len() {
+                    0 => None,
+                    1 => Some(Hover {
+                        contents: HoverContents::Scalar(comments[0].clone()),
+                        range: None,
+                    }),
+                    2.. => Some(Hover {
+                        contents: HoverContents::Array(comments),
+                        range: None,
+                    }),
+                };
+
+                Box::pin(async move { Ok(response) })
+            }
+        }
     }
 
     fn definition(
@@ -100,34 +87,20 @@ impl LanguageServer for ServerState {
         let uri = param.text_document_position_params.text_document.uri;
         let pos = param.text_document_position_params.position;
 
-        let Some(contents) = self.documents.get(&uri) else {
-            return Box::pin(async move {
-                Err(ResponseError::new(
-                    ErrorCode::INVALID_REQUEST,
-                    "uri was never opened",
-                ))
-            });
-        };
-
-        let Some(parsed) = self.parser.parse(contents.as_bytes()) else {
-            return Box::pin(async move {
-                Err(ResponseError::new(
-                    ErrorCode::REQUEST_FAILED,
-                    "ts failed to parse contents",
-                ))
-            });
-        };
-
-        let locations = parsed.definition(&pos, &uri, contents.as_bytes());
-        info!("Found {} matching nodes in the document", locations.len());
+        match self.get_parsed_tree_and_content(&uri) {
+            Err(e) => Box::pin(async move { Err(e) }),
+            Ok((tree, content)) => {
+                let locations = tree.definition(&pos, &uri, content.as_bytes());
 
-        let response = match locations.len() {
-            0 => None,
-            1 => Some(GotoDefinitionResponse::Scalar(locations[0].clone())),
-            2.. => Some(GotoDefinitionResponse::Array(locations)),
-        };
+                let response = match locations.len() {
+                    0 => None,
+                    1 => Some(GotoDefinitionResponse::Scalar(locations[0].clone())),
+                    2.. => Some(GotoDefinitionResponse::Array(locations)),
+                };
 
-        Box::pin(async move { Ok(response) })
+                Box::pin(async move { Ok(response) })
+            }
+        }
     }
 
     fn did_save(&mut self, _: DidSaveTextDocumentParams) -> Self::NotifyResult {
@@ -137,34 +110,45 @@ impl LanguageServer for ServerState {
     fn did_open(&mut self, params: DidOpenTextDocumentParams) -> Self::NotifyResult {
         let uri = params.text_document.uri;
         let contents = params.text_document.text;
-        info!("Opened file at: {:}", uri);
+
+        info!("opened file at: {uri}");
         self.documents.insert(uri.clone(), contents.clone());
 
-        let Some(parsed) = self.parser.parse(contents.as_bytes()) else {
-            tracing::error!("failed to parse content");
+        let Some(tree) = self.parser.parse(contents.as_bytes()) else {
+            error!("failed to parse content at {uri}");
             return ControlFlow::Continue(());
         };
 
-        let diagnostics = parsed.collect_parse_errors(&uri);
+        let diagnostics = tree.collect_parse_errors(&uri);
         if let Err(e) = self.client.publish_diagnostics(diagnostics) {
-            tracing::error!("failed to publish diagnostics. {:?}", e)
+            error!(error=%e, "failed to publish diagnostics")
         }
         ControlFlow::Continue(())
     }
 
+    fn did_close(&mut self, params: DidCloseTextDocumentParams) -> Self::NotifyResult {
+        let uri = params.text_document.uri;
+
+        info!("closed file at {uri}");
+        self.documents.remove(&uri);
+
+        ControlFlow::Continue(())
+    }
+
     fn did_change(&mut self, params: DidChangeTextDocumentParams) -> Self::NotifyResult {
         let uri = params.text_document.uri;
         let contents = params.content_changes[0].text.clone();
+
         self.documents.insert(uri.clone(), contents.clone());
 
-        let Some(parsed) = self.parser.parse(contents.as_bytes()) else {
-            tracing::error!("failed to parse content");
+        let Some(tree) = self.parser.parse(contents.as_bytes()) else {
+            error!("failed to parse content at {uri}");
             return ControlFlow::Continue(());
         };
 
-        let diagnostics = parsed.collect_parse_errors(&uri);
+        let diagnostics = tree.collect_parse_errors(&uri);
         if let Err(e) = self.client.publish_diagnostics(diagnostics) {
-            tracing::error!("failed to publish diagnostics. {:?}", e)
+            error!(error=%e, "failed to publish diagnostics")
         }
         ControlFlow::Continue(())
     }
@@ -175,28 +159,14 @@ impl LanguageServer for ServerState {
     ) -> BoxFuture<'static, Result<Option<DocumentSymbolResponse>, Self::Error>> {
         let uri = params.text_document.uri;
 
-        let Some(contents) = self.documents.get(&uri) else {
-            return Box::pin(async move {
-                Err(ResponseError::new(
-                    ErrorCode::INVALID_REQUEST,
-                    "uri was never opened",
-                ))
-            });
-        };
-
-        let Some(parsed) = self.parser.parse(contents.as_bytes()) else {
-            return Box::pin(async move {
-                Err(ResponseError::new(
-                    ErrorCode::REQUEST_FAILED,
-                    "ts failed to parse contents",
-                ))
-            });
-        };
+        match self.get_parsed_tree_and_content(&uri) {
+            Err(e) => Box::pin(async move { Err(e) }),
+            Ok((tree, content)) => {
+                let locations = tree.find_document_locations(content.as_bytes());
+                let response = DocumentSymbolResponse::Nested(locations);
 
-        let locations = parsed.find_document_locations(contents.as_bytes());
-
-        let response = DocumentSymbolResponse::Nested(locations);
-
-        Box::pin(async move { Ok(Some(response)) })
+                Box::pin(async move { Ok(Some(response)) })
+            }
+        }
     }
 }

src/lsp.rs

@@ -24,6 +24,7 @@ struct DocumentSymbolTreeBuilder {
     // The found are things we've finished processing/parsing, at the top level of the stack.
     found: Vec<DocumentSymbol>,
 }
+
 impl DocumentSymbolTreeBuilder {
     fn push(&mut self, node: usize, symbol: DocumentSymbol) {
         self.stack.push((node, symbol));
@@ -165,7 +166,6 @@ impl ParsedTree {
         let content = content.as_ref();
 
         let mut cursor = self.tree.root_node().walk();
-
         self.find_document_locations_inner(&mut builder, &mut cursor, content);
 
         builder.build()
@@ -191,6 +191,8 @@ impl ParsedTree {
                 let detail = self.find_preceding_comments(node.id(), content);
                 let message = node.parent().unwrap();
 
+                // https://github.com/rust-lang/rust/issues/102777
+                #[allow(deprecated)]
                 let new_symbol = DocumentSymbol {
                     name: name.to_string(),
                     detail,
@@ -252,6 +254,7 @@ impl ParsedTree {
     pub fn hover(&self, pos: &Position, content: impl AsRef<[u8]>) -> Vec<MarkedString> {
         let text = self.get_node_text_at_position(pos, content.as_ref());
         info!("Looking for hover response on: {:?}", text);
+
         match text {
             Some(text) => self
                 .find_childrens_by_kinds(&["message_name", "enum_name", "service_name", "rpc_name"])
@@ -427,6 +430,7 @@ Author has a name and a country where they were born"#
     }
 
     #[test]
+    #[allow(deprecated)]
     fn test_document_symbols() {
         let contents = r#"syntax = "proto3";
 

src/parser.rs

@@ -1,7 +1,8 @@
-use async_lsp::{lsp_types::Url, router::Router, ClientSocket};
+use async_lsp::{lsp_types::Url, router::Router, ClientSocket, ErrorCode, ResponseError};
 use std::{collections::HashMap, ops::ControlFlow};
+use tracing::error;
 
-use crate::parser::ProtoParser;
+use crate::parser::{ParsedTree, ProtoParser};
 
 pub struct TickEvent;
 pub struct ServerState {
@@ -27,4 +28,27 @@ impl ServerState {
         self.counter += 1;
         ControlFlow::Continue(())
     }
+
+    pub fn get_parsed_tree_and_content(
+        &mut self,
+        uri: &Url,
+    ) -> Result<(ParsedTree, &str), ResponseError> {
+        let Some(content) = self.documents.get(uri) else {
+            error!("failed to get document at {uri}");
+            return Err(ResponseError::new(
+                ErrorCode::INVALID_REQUEST,
+                "uri was never opened",
+            ));
+        };
+
+        let Some(parsed) = self.parser.parse(content.as_bytes()) else {
+            error!("failed to parse content at {uri}");
+            return Err(ResponseError::new(
+                ErrorCode::REQUEST_FAILED,
+                "ts failed to parse contents",
+            ));
+        };
+
+        Ok((parsed, content))
+    }
 }