Fix goto def

3 files changed, 68 insertions, 21 deletions

diff --git a/src/main/java/org/javacs/CompileFile.java b/src/main/java/org/javacs/CompileFile.java
index 297e3f5..826f7e9 100644
--- a/src/main/java/org/javacs/CompileFile.java
+++ b/src/main/java/org/javacs/CompileFile.java
@@ -6,8 +6,10 @@ import java.io.IOException;
 import java.net.URI;
 import java.util.HashSet;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.logging.Logger;
 import javax.lang.model.element.*;
+import org.eclipse.lsp4j.Range;
 
 public class CompileFile {
     private final JavaCompilerService parent;
@@ -37,6 +39,52 @@ public class CompileFile {
         profiler.print();
     }
 
+    public Optional<TreePath> find(Ptr target) {
+        class FindPtr extends TreePathScanner<Void, Void> {
+            TreePath found = null;
+
+            boolean toStringEquals(Object left, Object right) {
+                return Objects.equals(Objects.toString(left, ""), Objects.toString(right, ""));
+            }
+
+            /** Check if the declaration at the current path is the same symbol as `e` */
+            boolean sameSymbol() {
+                return new Ptr(getCurrentPath()).equals(target);
+            }
+
+            void check() {
+                if (sameSymbol()) {
+                    found = getCurrentPath();
+                }
+            }
+
+            @Override
+            public Void visitClass(ClassTree node, Void aVoid) {
+                check();
+                return super.visitClass(node, aVoid);
+            }
+
+            @Override
+            public Void visitMethod(MethodTree node, Void aVoid) {
+                check();
+                return super.visitMethod(node, aVoid);
+            }
+
+            @Override
+            public Void visitVariable(VariableTree node, Void aVoid) {
+                check();
+                return super.visitVariable(node, aVoid);
+            }
+        }
+        var find = new FindPtr();
+        find.scan(root, null);
+        return Optional.ofNullable(find.found);
+    }
+
+    public Optional<Range> range(TreePath path) {
+        return ParseFile.range(task, contents, path);
+    }
+
     /**
      * Figure out what imports this file should have. Star-imports like `import java.util.*` are converted to individual
      * class imports. Missing imports are inferred by looking at imports in other source files.
diff --git a/src/main/java/org/javacs/CompileFocus.java b/src/main/java/org/javacs/CompileFocus.java
index 48daa8a..c90e99b 100644
--- a/src/main/java/org/javacs/CompileFocus.java
+++ b/src/main/java/org/javacs/CompileFocus.java
@@ -70,14 +70,10 @@ public class CompileFocus {
                         Collections.singletonList(new StringFileObject(contents, file)));
     }
 
-    public Optional<TreePath> definition() {
-        LOG.info(String.format("Go-to-definition at %s(%d,%d)...", file, line, character));
-        var e = trees.getElement(path);
-        LOG.info(String.format("...found element `%s`", e));
-        var p = trees.getPath(e);
-        if (p != null) LOG.info(String.format("...element is located at path %s", new Ptr(p)));
-        else LOG.info(String.format("...element could not be located"));
-        return Optional.ofNullable(p);
+    public Optional<URI> declaringFile(Element e) {
+        var top = topLevelDeclaration(e);
+        if (!top.isPresent()) return Optional.empty();
+        return findDeclaringFile(top.get());
     }
 
     /** Find the smallest element that includes the cursor */
@@ -776,7 +772,7 @@ public class CompileFocus {
     }
 
     /** Find the file `e` was declared in */
-    private Optional<Path> findDeclaringFile(TypeElement e) {
+    private Optional<URI> findDeclaringFile(TypeElement e) {
         var name = e.getQualifiedName().toString();
         var lastDot = name.lastIndexOf('.');
         var packageName = lastDot == -1 ? "" : name.substring(0, lastDot);
@@ -787,7 +783,7 @@ public class CompileFocus {
         for (var root : parent.sourcePath) {
             var absPath = root.resolve(publicClassPath);
             if (Files.exists(absPath) && containsTopLevelDeclaration(absPath, className)) {
-                return Optional.of(absPath);
+                return Optional.of(absPath.toUri());
             }
         }
         // Then, look for a secondary declaration in all java files in the package
@@ -798,7 +794,7 @@ public class CompileFocus {
                 try {
                     var foundFile =
                             Files.list(absDir).filter(f -> containsTopLevelDeclaration(f, className)).findFirst();
-                    if (foundFile.isPresent()) return foundFile;
+                    if (foundFile.isPresent()) return foundFile.map(Path::toUri);
                 } catch (IOException err) {
                     throw new RuntimeException(err);
                 }
diff --git a/src/main/java/org/javacs/JavaTextDocumentService.java b/src/main/java/org/javacs/JavaTextDocumentService.java
index 3e80387..298b7c3 100644
--- a/src/main/java/org/javacs/JavaTextDocumentService.java
+++ b/src/main/java/org/javacs/JavaTextDocumentService.java
@@ -391,18 +391,18 @@ class JavaTextDocumentService implements TextDocumentService {
         var fromLine = position.getPosition().getLine() + 1;
         var fromColumn = position.getPosition().getCharacter() + 1;
         var fromContent = contents(fromUri).content;
-        var toPath = server.compiler.compileFocus(fromUri, fromContent, fromLine, fromColumn).definition();
+        var fromFocus = server.compiler.compileFocus(fromUri, fromContent, fromLine, fromColumn);
+        var toEl = fromFocus.element();
+        var toUri = fromFocus.declaringFile(toEl);
+        if (!toUri.isPresent()) return CompletableFuture.completedFuture(List.of());
+        var toContent = contents(toUri.get()).content;
+        var toFile = server.compiler.compileFile(toUri.get(), toContent);
+        var toPath = toFile.find(new Ptr(toEl));
         if (!toPath.isPresent()) return CompletableFuture.completedFuture(List.of());
-        // Figure out what file definition is in
-        var toUri = toPath.get().getCompilationUnit().getSourceFile().toUri();
-        // Parse that file
-        var toContent = contents(toUri).content;
-        var parse = server.compiler.parseFile(toUri, toContent);
         // Figure out where in the file the definition is
-        // NOTE: toPath is from a different compilation than parse, but this seems to work
-        var toRange = parse.range(toPath.get());
+        var toRange = toFile.range(toPath.get());
         if (!toRange.isPresent()) return CompletableFuture.completedFuture(List.of());
-        var to = new Location(toUri.toString(), toRange.get());
+        var to = new Location(toUri.get().toString(), toRange.get());
         return CompletableFuture.completedFuture(List.of(to));
     }
 
@@ -525,7 +525,10 @@ class JavaTextDocumentService implements TextDocumentService {
                 // TODO run all tests in package
             } else if (parse.isTestMethod(d)) {
                 var command =
-                        new Command("Run Test", "java.command.test.run", Arrays.asList(uri, className, memberName));
+                        new Command(
+                                "Run Test",
+                                "java.command.test.run",
+                                Arrays.asList(uri, className, memberName.orElse(null)));
                 var lens = new CodeLens(range.get(), command, null);
                 result.add(lens);
             }