feat: recursively parse import files and goto definition of well know…

…n types (#56) When a file is opened, all imports are parsed recursively up to a depth of 8. When a file is being edited it and its imports are parsed with a depth of 2.
coder3101 · Jan 26, 2025 · adda8ae · adda8ae
1 parent f392094
commit adda8ae
Show file tree

Hide file tree

Showing 18 changed files with 150 additions and 61 deletions.
diff --git a/README.md b/README.md
@@ -15,7 +15,7 @@
 - ✅ **Diagnostics**: Syntax errors and import error detected with the tree-sitter parser.
 - ✅ **Document Symbols**: Navigate and view all symbols, including messages and enums.
 - ✅ **Code Formatting**: Format `.proto` files using `clang-format` for a consistent style.
-- ✅ **Go to Definition**: Jump to the definition of symbols like messages or enums.
+- ✅ **Go to Definition**: Jump to the definition of symbols like messages or enums and imports.
 - ✅ **Hover Information**: Get detailed information and documentation on hover.
 - ✅ **Rename Symbols**: Rename protobuf symbols and propagate changes across the codebase.
 - ✅ **Find References**: Find where messages, enums, and fields are used throughout the codebase.
@@ -134,7 +134,7 @@ Protols provides a list of symbols in the current document, including nested sym
 
 ### Go to Definition
 
-Jump directly to the definition of any custom symbol, including those in other files or packages. This feature works across package boundaries.
+Jump directly to the definition of any custom symbol or imports, including those in other files or packages. This feature works across package boundaries.
 
 ### Hover Information
 

diff --git a/src/lsp.rs b/src/lsp.rs
@@ -381,7 +381,7 @@ impl LanguageServer for ProtoLanguageServer {
             return ControlFlow::Continue(());
         };
 
-        let Some(diagnostics) = self.state.upsert_file(&uri, content.clone(), &ipath) else {
+        let Some(diagnostics) = self.state.upsert_file(&uri, content.clone(), &ipath, 8) else {
             return ControlFlow::Continue(());
         };
 
@@ -405,7 +405,7 @@ impl LanguageServer for ProtoLanguageServer {
             return ControlFlow::Continue(());
         };
 
-        let Some(diagnostics) = self.state.upsert_file(&uri, content, &ipath) else {
+        let Some(diagnostics) = self.state.upsert_file(&uri, content, &ipath, 2) else {
             return ControlFlow::Continue(());
         };
 
@@ -427,9 +427,10 @@ impl LanguageServer for ProtoLanguageServer {
             if let Ok(uri) = Url::from_file_path(&file.uri) {
                 // Safety: The uri is always a file type
                 let content = read_to_string(uri.to_file_path().unwrap()).unwrap_or_default();
-                self.state.upsert_content(&uri, content, &[]);
-            } else {
-                error!(uri=%file.uri, "failed parse uri");
+
+                if let Some(ipath) = self.configs.get_include_paths(&uri) {
+                    self.state.upsert_content(&uri, content, &ipath, 2);
+                }
             }
         }
         ControlFlow::Continue(())

diff --git a/src/state.rs b/src/state.rs
@@ -1,7 +1,7 @@
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     path::PathBuf,
-    sync::{Arc, Mutex, MutexGuard, RwLock, RwLockWriteGuard},
+    sync::{Arc, Mutex, RwLock},
 };
 use tracing::info;
 
@@ -66,34 +66,74 @@ impl ProtoLanguageState {
     }
 
     fn upsert_content_impl(
-        mut parser: MutexGuard<ProtoParser>,
+        &mut self,
         uri: &Url,
         content: String,
-        mut docs: RwLockWriteGuard<HashMap<Url, String>>,
-        mut trees: RwLockWriteGuard<HashMap<Url, ParsedTree>>,
-    ) -> bool {
-        if let Some(parsed) = parser.parse(uri.clone(), content.as_bytes()) {
-            trees.insert(uri.clone(), parsed);
-            docs.insert(uri.clone(), content);
-            true
-        } else {
-            false
+        ipath: &[PathBuf],
+        depth: usize,
+        parse_session: &mut HashSet<Url>,
+    ) {
+        // Safety: to not cause stack overflow
+        if depth == 0 {
+            return;
+        }
+
+        // avoid re-parsing same file incase of circular dependencies
+        if parse_session.contains(uri) {
+            return;
         }
+
+        let Some(parsed) = self
+            .parser
+            .lock()
+            .expect("poison")
+            .parse(uri.clone(), content.as_bytes())
+        else {
+            return;
+        };
+
+        self.trees
+            .write()
+            .expect("posion")
+            .insert(uri.clone(), parsed);
+
+        self.documents
+            .write()
+            .expect("poison")
+            .insert(uri.clone(), content.clone());
+
+        parse_session.insert(uri.clone());
+        let imports = self.get_owned_imports(uri, content.as_str());
+
+        for import in imports.iter() {
+            if let Some(p) = ipath.iter().map(|p| p.join(import)).find(|p| p.exists()) {
+                if let Ok(uri) = Url::from_file_path(p.clone()) {
+                    if let Ok(content) = std::fs::read_to_string(p) {
+                        self.upsert_content_impl(&uri, content, ipath, depth - 1, parse_session);
+                    }
+                }
+            }
+        }
+    }
+
+    fn get_owned_imports(&self, uri: &Url, content: &str) -> Vec<String> {
+        self.get_tree(uri)
+            .map(|t| t.get_import_paths(content.as_ref()))
+            .unwrap_or_default()
+            .into_iter()
+            .map(ToOwned::to_owned)
+            .collect()
     }
 
     pub fn upsert_content(
         &mut self,
         uri: &Url,
         content: String,
         ipath: &[PathBuf],
+        depth: usize,
     ) -> Vec<String> {
-        // Drop locks at end of block
-        {
-            let parser = self.parser.lock().expect("poison");
-            let tree = self.trees.write().expect("poison");
-            let docs = self.documents.write().expect("poison");
-            Self::upsert_content_impl(parser, uri, content.clone(), docs, tree);
-        }
+        let mut session = HashSet::new();
+        self.upsert_content_impl(uri, content.clone(), ipath, depth, &mut session);
 
         // After content is upserted, those imports which couldn't be located
         // are flagged as import error
@@ -189,9 +229,10 @@ impl ProtoLanguageState {
         uri: &Url,
         content: String,
         ipath: &[PathBuf],
+        depth: usize,
     ) -> Option<PublishDiagnosticsParams> {
-        info!(uri=%uri, "upserting file");
-        let diag = self.upsert_content(uri, content.clone(), ipath);
+        info!(%uri, %depth, "upserting file");
+        let diag = self.upsert_content(uri, content.clone(), ipath, depth);
         self.get_tree(uri).map(|tree| {
             let diag = tree.collect_import_diagnostics(content.as_ref(), diag);
             let mut d = tree.collect_parse_diagnostics();
@@ -205,13 +246,13 @@ impl ProtoLanguageState {
     }
 
     pub fn delete_file(&mut self, uri: &Url) {
-        info!(uri=%uri, "deleting file");
+        info!(%uri, "deleting file");
         self.documents.write().expect("poison").remove(uri);
         self.trees.write().expect("poison").remove(uri);
     }
 
     pub fn rename_file(&mut self, new_uri: &Url, old_uri: &Url) {
-        info!(new_uri=%new_uri, old_uri=%new_uri, "renaming file");
+        info!(%new_uri, %new_uri, "renaming file");
 
         if let Some(v) = self.documents.write().expect("poison").remove(old_uri) {
             self.documents

diff --git a/src/workspace/definition.rs b/src/workspace/definition.rs
@@ -66,9 +66,9 @@ mod test {
         let c = include_str!("input/c.proto");
 
         let mut state: ProtoLanguageState = ProtoLanguageState::new();
-        state.upsert_file(&a_uri, a.to_owned(), &ipath);
-        state.upsert_file(&b_uri, b.to_owned(), &ipath);
-        state.upsert_file(&c_uri, c.to_owned(), &ipath);
+        state.upsert_file(&a_uri, a.to_owned(), &ipath, 2);
+        state.upsert_file(&b_uri, b.to_owned(), &ipath, 2);
+        state.upsert_file(&c_uri, c.to_owned(), &ipath, 2);
 
         assert_yaml_snapshot!(state.definition(
             &ipath,

diff --git a/src/workspace/hover.rs b/src/workspace/hover.rs
@@ -672,7 +672,7 @@ mod test {
 
     #[test]
     fn workspace_test_hover() {
-        let ipath = vec![PathBuf::from("src/workspace/input")];
+        let ipath = vec![std::env::current_dir().unwrap().join("src/workspace/input")];
         let a_uri = "file://input/a.proto".parse().unwrap();
         let b_uri = "file://input/b.proto".parse().unwrap();
         let c_uri = "file://input/c.proto".parse().unwrap();
@@ -682,9 +682,9 @@ mod test {
         let c = include_str!("input/c.proto");
 
         let mut state: ProtoLanguageState = ProtoLanguageState::new();
-        state.upsert_file(&a_uri, a.to_owned(), &ipath);
-        state.upsert_file(&b_uri, b.to_owned(), &ipath);
-        state.upsert_file(&c_uri, c.to_owned(), &ipath);
+        state.upsert_file(&a_uri, a.to_owned(), &ipath, 3);
+        state.upsert_file(&b_uri, b.to_owned(), &ipath, 2);
+        state.upsert_file(&c_uri, c.to_owned(), &ipath, 2);
 
         assert_yaml_snapshot!(state.hover(
             &ipath,
@@ -719,7 +719,12 @@ mod test {
         assert_yaml_snapshot!(state.hover(
             &ipath,
             "com.workspace",
-            Hoverables::ImportPath("c.proto".to_string())
-        ))
+            Hoverables::Identifier("com.inner.Why".to_string())
+        ));
+        assert_yaml_snapshot!(state.hover(
+            &ipath,
+            "com.workspace",
+            Hoverables::Identifier("com.super.secret.SomeSecret".to_string())
+        ));
     }
 }
diff --git a/src/workspace/input/a.proto b/src/workspace/input/a.proto
@@ -3,11 +3,13 @@ syntax = "proto3";
 package com.workspace;
 
 import "c.proto";
+import "inner/x.proto";
 
 // A Book is a book
 message Book {
    Author author = 1;
    Author.Address foo = 2;
    com.utility.Foobar.Baz z = 3;
+   com.inner.Why reason = 4;
 }
 
diff --git a/src/workspace/input/inner/secret/y.proto b/src/workspace/input/inner/secret/y.proto
@@ -0,0 +1,8 @@
+syntax = "proto3";
+
+package com.super.secret;
+
+// SomeSecret is a real secret with hidden string
+message SomeSecret {
+   string secret = 1;
+}
diff --git a/src/workspace/input/inner/x.proto b/src/workspace/input/inner/x.proto
@@ -0,0 +1,12 @@
+syntax = "proto3";
+
+package com.inner;
+
+import "inner/secret/y.proto";
+
+// Why is a reason with secret
+message Why {
+   string reason = 1;
+   com.super.secret.SomeSecret secret = 2;
+}
+
diff --git a/src/workspace/rename.rs b/src/workspace/rename.rs
@@ -79,9 +79,9 @@ mod test {
         let c = include_str!("input/c.proto");
 
         let mut state: ProtoLanguageState = ProtoLanguageState::new();
-        state.upsert_file(&a_uri, a.to_owned(), &ipath);
-        state.upsert_file(&b_uri, b.to_owned(), &ipath);
-        state.upsert_file(&c_uri, c.to_owned(), &ipath);
+        state.upsert_file(&a_uri, a.to_owned(), &ipath, 2);
+        state.upsert_file(&b_uri, b.to_owned(), &ipath, 2);
+        state.upsert_file(&c_uri, c.to_owned(), &ipath, 2);
 
         assert_yaml_snapshot!(state.rename_fields("com.workspace", "Author", "Writer"));
         assert_yaml_snapshot!(state.rename_fields(
@@ -104,9 +104,9 @@ mod test {
         let c = include_str!("input/c.proto");
 
         let mut state: ProtoLanguageState = ProtoLanguageState::new();
-        state.upsert_file(&a_uri, a.to_owned(), &ipath);
-        state.upsert_file(&b_uri, b.to_owned(), &ipath);
-        state.upsert_file(&c_uri, c.to_owned(), &ipath);
+        state.upsert_file(&a_uri, a.to_owned(), &ipath, 2);
+        state.upsert_file(&b_uri, b.to_owned(), &ipath, 2);
+        state.upsert_file(&c_uri, c.to_owned(), &ipath, 2);
 
         assert_yaml_snapshot!(state.reference_fields("com.workspace", "Author"));
         assert_yaml_snapshot!(state.reference_fields("com.workspace", "Author.Address"));

diff --git a/src/workspace/snapshots/protols__workspace__hover__test__workspace_test_hover-7.snap b/src/workspace/snapshots/protols__workspace__hover__test__workspace_test_hover-7.snap
@@ -1,7 +1,7 @@
 ---
 source: src/workspace/hover.rs
-expression: "state.hover(&ipath, \"com.workspace\",\nHoverables::ImportPath(\"c.proto\".to_string()))"
+expression: "state.hover(&ipath, \"com.workspace\",\nHoverables::Identifier(\"com.inner.Why\".to_string()))"
 snapshot_kind: text
 ---
 kind: markdown
-value: "Import: `c.proto` protobuf file,\n---\nIncluded from src/workspace/input/c.proto"
+value: "`Why` message or enum type, package: `com.inner`\n---\nWhy is a reason with secret"
diff --git a/src/workspace/snapshots/protols__workspace__hover__test__workspace_test_hover-8.snap b/src/workspace/snapshots/protols__workspace__hover__test__workspace_test_hover-8.snap
@@ -0,0 +1,7 @@
+---
+source: src/workspace/hover.rs
+expression: "state.hover(&ipath, \"com.workspace\",\nHoverables::Identifier(\"com.super.secret.SomeSecret\".to_string()))"
+snapshot_kind: text
+---
+kind: markdown
+value: "`SomeSecret` message or enum type, package: `com.super.secret`\n---\nSomeSecret is a real secret with hidden string"
diff --git a/src/workspace/snapshots/protols__workspace__hover__test__workspace_test_hover-9.snap b/src/workspace/snapshots/protols__workspace__hover__test__workspace_test_hover-9.snap
@@ -0,0 +1,7 @@
+---
+source: src/workspace/hover.rs
+expression: "state.hover(&ipath, \"com.workspace\",\nHoverables::Identifier(\"com.super.secret.SomeSecret\".to_string()))"
+snapshot_kind: text
+---
+kind: markdown
+value: "`SomeSecret` message or enum type, package: `com.super.secret`\n---\nSomeSecret is a real secret with hidden string"
diff --git a/src/workspace/snapshots/protols__workspace__rename__test__reference-2.snap b/src/workspace/snapshots/protols__workspace__rename__test__reference-2.snap
@@ -1,12 +1,13 @@
 ---
 source: src/workspace/rename.rs
 expression: "state.reference_fields(\"com.workspace\", \"Author.Address\")"
+snapshot_kind: text
 ---
 - uri: "file://input/a.proto"
   range:
     start:
-      line: 9
+      line: 10
       character: 3
     end:
-      line: 9
+      line: 10
       character: 17
diff --git a/src/workspace/snapshots/protols__workspace__rename__test__reference-3.snap b/src/workspace/snapshots/protols__workspace__rename__test__reference-3.snap
@@ -1,12 +1,13 @@
 ---
 source: src/workspace/rename.rs
 expression: "state.reference_fields(\"com.utility\", \"Foobar.Baz\")"
+snapshot_kind: text
 ---
 - uri: "file://input/a.proto"
   range:
     start:
-      line: 10
+      line: 11
       character: 3
     end:
-      line: 10
+      line: 11
       character: 25
diff --git a/src/workspace/snapshots/protols__workspace__rename__test__reference.snap b/src/workspace/snapshots/protols__workspace__rename__test__reference.snap
@@ -1,12 +1,13 @@
 ---
 source: src/workspace/rename.rs
 expression: "state.reference_fields(\"com.workspace\", \"Author\")"
+snapshot_kind: text
 ---
 - uri: "file://input/a.proto"
   range:
     start:
-      line: 8
+      line: 9
       character: 3
     end:
-      line: 8
+      line: 9
       character: 9
diff --git a/src/workspace/snapshots/protols__workspace__rename__test__rename-2.snap b/src/workspace/snapshots/protols__workspace__rename__test__rename-2.snap
@@ -1,13 +1,14 @@
 ---
 source: src/workspace/rename.rs
 expression: "state.rename_fields(\"com.workspace\", \"Author.Address\", \"Author.Location\")"
+snapshot_kind: text
 ---
 "file://input/a.proto":
   - range:
       start:
-        line: 9
+        line: 10
         character: 3
       end:
-        line: 9
+        line: 10
         character: 17
     newText: Author.Location