From 9ac22e59bccf7eb3c96380db4d702dae75929b74 Mon Sep 17 00:00:00 2001
From: Ryan Roden-Corrent <git@rrc.codes>
Date: Sat, 2 Nov 2024 15:04:09 -0400
Subject: [PATCH] Improve goto definition across packages.

---
 src/workspace.rs           | 74 ++++++++++++++++++++++++++++++++------
 testdata/folder/what.proto |  3 +-
 tests/integration_test.rs  | 30 ++++++++++++++++
 3 files changed, 95 insertions(+), 12 deletions(-)

diff --git a/src/workspace.rs b/src/workspace.rs
index 1f7c217..56abb43 100644
--- a/src/workspace.rs
+++ b/src/workspace.rs
@@ -37,6 +37,54 @@ pub struct Workspace {
     files: std::collections::HashMap<Url, file::File>,
 }
 
+// Return the possible package qualifiers to_pkg could use for a type imported from from_pkg
+fn possible_qualifiers<'a>(from_pkg: &'a str, to_pkg: &'a str) -> Vec<&'a str> {
+    log::trace!("possible_qualifiers({from_pkg}, {to_pkg})");
+    if to_pkg == "" {
+        return vec![from_pkg];
+    }
+
+    let mut res = vec![];
+    if let Some(pkg) = from_pkg.strip_prefix(to_pkg) {
+        let pkg = pkg.strip_prefix(".").unwrap_or(pkg);
+        res.push(pkg);
+    }
+
+    if let Some((to_pkg, _)) = to_pkg.rsplit_once(".") {
+        res.append(&mut possible_qualifiers(from_pkg, to_pkg));
+    } else {
+        res.push(from_pkg);
+    }
+    return res;
+}
+
+#[test]
+fn test_possible_qualifiers() {
+    let _ = env_logger::builder().is_test(true).try_init();
+    assert_eq!(possible_qualifiers("", ""), vec![""]);
+    assert_eq!(possible_qualifiers("foo", ""), vec!["foo"]);
+    assert_eq!(possible_qualifiers("foo.bar", ""), vec!["foo.bar"]);
+    assert_eq!(possible_qualifiers("foo", "bar"), vec!["foo"]);
+    assert_eq!(possible_qualifiers("foo.bar", "bar"), vec!["foo.bar"]);
+    assert_eq!(possible_qualifiers("foo", "foo"), vec!["", "foo"]);
+    assert_eq!(
+        possible_qualifiers("foo.bar", "foo"),
+        vec!["bar", "foo.bar"]
+    );
+    assert_eq!(
+        possible_qualifiers("foo.bar.baz", "foo.bar"),
+        vec!["baz", "bar.baz", "foo.bar.baz",]
+    );
+    assert_eq!(
+        possible_qualifiers("foo.bar.baz", "foo.bar.baz"),
+        vec!["", "baz", "bar.baz", "foo.bar.baz",]
+    );
+    assert_eq!(
+        possible_qualifiers("folder.stuff", "folder.what"),
+        vec!["stuff", "folder.stuff"]
+    );
+}
+
 impl Workspace {
     pub fn new(proto_paths: Vec<std::path::PathBuf>) -> Workspace {
         Workspace {
@@ -362,21 +410,25 @@ impl Workspace {
         for (uri, file) in imports {
             let package = file.package();
             if let Some(sym) = if package == local_package {
-                log::trace!("Searching for {} in {uri}", typ.name);
+                log::trace!("Searching for {} in {uri} (same package)", typ.name);
                 // same package, match the name without the package prefix
                 file.symbols(&mut qc).find(|sym| sym.name == typ.name)
             } else if let Some(package) = package {
+                log::trace!("Searching for {} in {uri} (different package)", typ.name);
                 // different package, fully qualify the name
-                log::trace!("Stripping {package} from {}", typ.name);
-                let qualified = typ
-                    .name
-                    .strip_prefix(package)
-                    .unwrap_or(typ.name)
-                    .strip_prefix(".")
-                    .unwrap_or(typ.name)
-                    .to_string();
-                log::trace!("Searching for {} in {uri} (different package)", qualified);
-                file.symbols(&mut qc).find(|sym| sym.name == qualified)
+                let local_package = local_package.unwrap_or("");
+                file.symbols(&mut qc).find(|sym| {
+                    let quals = possible_qualifiers(package, local_package);
+                    log::trace!("Qualifiers: {quals:?}");
+                    quals
+                        .iter()
+                        .inspect(|q| log::trace!("Qual == {q}"))
+                        .filter_map(|qual| typ.name.strip_prefix(qual))
+                        .inspect(|q| log::trace!("stripped == {q}"))
+                        .filter_map(|name| name.strip_prefix("."))
+                        .inspect(|q| log::trace!("name == {q} == {}", sym.name))
+                        .any(|name| name == sym.name)
+                })
             } else {
                 // target file has no package
                 log::trace!("Searching for {} in {uri}", typ.name);
diff --git a/testdata/folder/what.proto b/testdata/folder/what.proto
index 9f96969..cf67e23 100644
--- a/testdata/folder/what.proto
+++ b/testdata/folder/what.proto
@@ -5,6 +5,7 @@ package folder.what;
 import "folder/stuff.proto";
 
 message What{
-    stuff.Stuff stuff = 1;
+    stuff.Stuff a = 1;
+    folder.stuff.Stuff b = 2;
 }
 
diff --git a/tests/integration_test.rs b/tests/integration_test.rs
index 37bb405..c2e21e3 100644
--- a/tests/integration_test.rs
+++ b/tests/integration_test.rs
@@ -40,6 +40,10 @@ fn stuff_uri() -> Url {
     Url::from_file_path(std::fs::canonicalize("./testdata/folder/stuff.proto").unwrap()).unwrap()
 }
 
+fn what_uri() -> Url {
+    Url::from_file_path(std::fs::canonicalize("./testdata/folder/what.proto").unwrap()).unwrap()
+}
+
 fn diag(uri: Url, target: &str, message: &str) -> Diagnostic {
     Diagnostic {
         range: locate_sym(uri, target).range,
@@ -653,6 +657,32 @@ fn test_goto_definition_different_file_nested() -> pbls::Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_goto_definition_partially_qualified_package() -> pbls::Result<()> {
+    let mut client = TestClient::new()?;
+    client.open(what_uri())?;
+
+    let resp = client.request::<GotoDefinition>(goto(what_uri(), "stuff.Stuff a =", 0))?;
+    assert_eq!(
+        resp,
+        Some(GotoDefinitionResponse::Scalar(locate_sym(
+            stuff_uri(),
+            "message Stuff",
+        )))
+    );
+
+    let resp = client.request::<GotoDefinition>(goto(what_uri(), "folder.stuff.Stuff b =", 0))?;
+    assert_eq!(
+        resp,
+        Some(GotoDefinitionResponse::Scalar(locate_sym(
+            stuff_uri(),
+            "message Stuff",
+        )))
+    );
+
+    Ok(())
+}
+
 #[test]
 fn test_message_references() -> pbls::Result<()> {
     let mut client = TestClient::new()?;