feat: add completion for enum and message in a package

Author: coder3101

README.md

@@ -4,12 +4,12 @@ A Language Server for **proto3** files. It uses tree-sitter parser for all opera
 ![](./assets/protols.mov)
 
 ## Features 
-- [x] Hover 
-- [x] Go to definition
-- [x] Diagnostics
-- [x] Document Symbols for message and enums in current document
-- [x] Rename message, enum and rpc
-- [x] Completion for proto3 keywords
+- [x] Completion (keywords, enums and messages of the package)
+- [x] Diagnostics - based on sytax errors
+- [x] Document Symbols for message and enums
+- [x] Go to definition - across packages
+- [x] Hover - across packages
+- [x] Rename - in current buffer only
 
 ## Installation
 

sample/simple.proto

@@ -9,12 +9,18 @@ message Book {
     string title = 2;
     Author author = 3;
     google.protobuf.Any data = 4;
+    BookState state = 5;
 
     // Author is a author of a book
     message Author {
         string name = 1;
         int64 age = 2;
     }
+
+    enum BookState {
+        HARD_COVER = 1;
+        SOFT_COVER = 2;
+    }
 }
 
 // This is a comment on message

src/lsp.rs

@@ -156,14 +156,16 @@ impl LanguageServer for ProtoLanguageServer {
     }
     fn completion(
         &mut self,
-        _params: CompletionParams,
+        params: CompletionParams,
     ) -> BoxFuture<'static, Result<Option<CompletionResponse>, Self::Error>> {
+        let uri = params.text_document_position.text_document.uri;
+
         let keywords = vec![
             "syntax", "package", "option", "import", "service", "rpc", "returns", "message",
             "enum", "oneof", "repeated", "reserved", "to",
         ];
 
-        let keywords = keywords
+        let mut keywords: Vec<CompletionItem> = keywords
             .into_iter()
             .map(|w| CompletionItem {
                 label: w.to_string(),
@@ -172,6 +174,12 @@ impl LanguageServer for ProtoLanguageServer {
             })
             .collect();
 
+        if let Some(tree) = self.state.get_tree(&uri) {
+            let content = self.state.get_content(&uri);
+            if let Some(package_name) = tree.get_package_name(content.as_bytes()) {
+                keywords.extend(self.state.completion_items(package_name));
+            }
+        }
         Box::pin(async move { Ok(Some(CompletionResponse::Array(keywords))) })
     }
 

src/main.rs

@@ -10,6 +10,7 @@ use tower::ServiceBuilder;
 use tracing::Level;
 
 mod lsp;
+mod nodekind;
 mod parser;
 mod server;
 mod state;

src/parser/nodekind.rs renamed to src/nodekind.rs

@@ -39,6 +39,14 @@ impl NodeKind {
         n.kind() == Self::PackageName.as_str()
     }
 
+    pub fn is_enum_name(n: &Node) -> bool {
+        n.kind() == Self::EnumName.as_str()
+    }
+
+    pub fn is_message_name(n: &Node) -> bool {
+        n.kind() == Self::MessageName.as_str()
+    }
+
     pub fn is_userdefined(n: &Node) -> bool {
         n.kind() == Self::EnumName.as_str() || n.kind() == Self::MessageName.as_str()
     }

src/parser/definition.rs

@@ -1,7 +1,7 @@
 use async_lsp::lsp_types::{Location, Range};
 use tree_sitter::Node;
 
-use crate::{parser::nodekind::NodeKind, utils::ts_to_lsp_position};
+use crate::{nodekind::NodeKind, utils::ts_to_lsp_position};
 
 use super::ParsedTree;
 

src/parser/diagnostics.rs

@@ -1,8 +1,8 @@
 use async_lsp::lsp_types::{Diagnostic, DiagnosticSeverity, PublishDiagnosticsParams, Range};
 
-use crate::utils::ts_to_lsp_position;
+use crate::{nodekind::NodeKind, utils::ts_to_lsp_position};
 
-use super::{nodekind::NodeKind, ParsedTree};
+use super::ParsedTree;
 
 impl ParsedTree {
     pub fn collect_parse_errors(&self) -> PublishDiagnosticsParams {

src/parser/docsymbol.rs

@@ -1,9 +1,9 @@
 use async_lsp::lsp_types::{DocumentSymbol, Range};
 use tree_sitter::TreeCursor;
 
-use crate::utils::ts_to_lsp_position;
+use crate::{nodekind::NodeKind, utils::ts_to_lsp_position};
 
-use super::{nodekind::NodeKind, ParsedTree};
+use super::ParsedTree;
 
 #[derive(Default)]
 pub(super) struct DocumentSymbolTreeBuilder {

src/parser/hover.rs

@@ -1,7 +1,7 @@
 use async_lsp::lsp_types::MarkedString;
 use tree_sitter::Node;
 
-use crate::parser::nodekind::NodeKind;
+use crate::nodekind::NodeKind;
 
 use super::ParsedTree;
 

src/parser/input/test_can_rename.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package com.book;
+package com.parser;
 
 // A Book is book
 message Book {

src/parser/input/test_collect_parse_error1.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package test;
+package com.parser;
 
 message Foo {
 	reserved 1;

src/parser/input/test_collect_parse_error2.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package com.book;
+package com.parser;
 
 message Book {
     message Author {

src/parser/input/test_document_symbols.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package com.symbols;
+package com.parser;
 
 // outer 1 comment
 message Outer1 {

src/parser/input/test_filter.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package com.book;
+package com.parser;
 
 message Book {
 

src/parser/input/test_goto_definition.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package com.book;
+package com.parser;
 
 message Book {
     message Author {

src/parser/input/test_hover.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package com.book;
+package com.parser;
 
 // A Book is book
 message Book {

src/parser/input/test_rename.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package com.book;
+package com.parser;
 
 // A Book is book
 message Book {

src/parser/mod.rs

@@ -5,7 +5,6 @@ mod definition;
 mod diagnostics;
 mod docsymbol;
 mod hover;
-mod nodekind;
 mod rename;
 mod tree;
 

src/parser/rename.rs

@@ -2,9 +2,9 @@ use std::collections::HashMap;
 
 use async_lsp::lsp_types::{Position, Range, TextEdit, WorkspaceEdit};
 
-use crate::utils::ts_to_lsp_position;
+use crate::{nodekind::NodeKind, utils::ts_to_lsp_position};
 
-use super::{nodekind::NodeKind, ParsedTree};
+use super::ParsedTree;
 
 impl ParsedTree {
     pub fn can_rename(&self, pos: &Position) -> Option<Range> {

src/parser/snapshots/protols__parser__tree__test__filter-2.snap

@@ -2,4 +2,4 @@
 source: src/parser/tree.rs
 expression: package_name
 ---
-com.book
+com.parser

src/parser/tree.rs

@@ -1,9 +1,9 @@
 use async_lsp::lsp_types::Position;
 use tree_sitter::{Node, TreeCursor};
 
-use crate::utils::lsp_to_ts_point;
+use crate::{nodekind::NodeKind, utils::lsp_to_ts_point};
 
-use super::{nodekind::NodeKind, ParsedTree};
+use super::ParsedTree;
 
 impl ParsedTree {
     pub(super) fn walk_and_collect_filter<'a>(

src/state.rs

@@ -1,10 +1,16 @@
 use std::{collections::HashMap, fs::read_to_string};
 use tracing::{error, info};
 
-use async_lsp::lsp_types::{PublishDiagnosticsParams, Url, WorkspaceFolder};
+use async_lsp::lsp_types::{
+    CompletionItem, CompletionItemKind, PublishDiagnosticsParams, Url, WorkspaceFolder,
+};
+use tree_sitter::Node;
 use walkdir::WalkDir;
 
-use crate::parser::{ParsedTree, ProtoParser};
+use crate::{
+    nodekind::NodeKind,
+    parser::{ParsedTree, ProtoParser},
+};
 
 pub struct ProtoLanguageState {
     documents: HashMap<Url, String>,
@@ -106,4 +112,31 @@ impl ProtoLanguageState {
             self.trees.insert(new_uri.clone(), v);
         }
     }
+
+    pub fn completion_items(&self, package: &str) -> Vec<CompletionItem> {
+        let collector = |f: fn(&Node) -> bool, k: CompletionItemKind| {
+            self.get_trees_for_package(package)
+                .into_iter()
+                .fold(vec![], |mut v, tree| {
+                    let content = self.get_content(&tree.uri);
+                    let t = tree.filter_nodes(f).into_iter().map(|n| CompletionItem {
+                        label: n.utf8_text(content.as_bytes()).unwrap().to_string(),
+                        kind: Some(k),
+                        ..Default::default()
+                    });
+                    v.extend(t);
+                    return v;
+                })
+        };
+
+        let mut result = collector(NodeKind::is_enum_name, CompletionItemKind::ENUM);
+        result.extend(collector(
+            NodeKind::is_message_name,
+            CompletionItemKind::STRUCT,
+        ));
+        // Better ways to dedup, but who cares?...
+        result.sort_by_key(|k| k.label.clone());
+        result.dedup_by_key(|k| k.label.clone());
+        result
+    }
 }

src/workspace/definition.rs

@@ -38,9 +38,9 @@ mod test {
         state.upsert_file(&b_uri, b.to_owned());
         state.upsert_file(&c_uri, c.to_owned());
 
-        assert_yaml_snapshot!(state.definition("com.library", "Author"));
-        assert_yaml_snapshot!(state.definition("com.library", "Author.Address"));
-        assert_yaml_snapshot!(state.definition("com.library", "com.utility.Foobar.Baz"));
+        assert_yaml_snapshot!(state.definition("com.workspace", "Author"));
+        assert_yaml_snapshot!(state.definition("com.workspace", "Author.Address"));
+        assert_yaml_snapshot!(state.definition("com.workspace", "com.utility.Foobar.Baz"));
         assert_yaml_snapshot!(state.definition("com.utility", "Baz"));
     }
 }

src/workspace/hover.rs

@@ -39,9 +39,9 @@ mod test {
         state.upsert_file(&b_uri, b.to_owned());
         state.upsert_file(&c_uri, c.to_owned());
 
-        assert_yaml_snapshot!(state.hover("com.library", "Author"));
-        assert_yaml_snapshot!(state.hover("com.library", "Author.Address"));
-        assert_yaml_snapshot!(state.hover("com.library", "com.utility.Foobar.Baz"));
+        assert_yaml_snapshot!(state.hover("com.workspace", "Author"));
+        assert_yaml_snapshot!(state.hover("com.workspace", "Author.Address"));
+        assert_yaml_snapshot!(state.hover("com.workspace", "com.utility.Foobar.Baz"));
         assert_yaml_snapshot!(state.hover("com.utility", "Baz"));
     }
 }

src/workspace/input/a.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package com.library;
+package com.workspace;
 
 import "c.proto";
 

src/workspace/input/b.proto

@@ -1,6 +1,6 @@
 syntax = "proto3";
 
-package com.library;
+package com.workspace;
 
 // A author is a author
 message Author {
@@ -13,4 +13,3 @@ message Author {
 
    Address foo = 2;
 }
-