Reduce use of TreePath

3 files changed, 85 insertions, 122 deletions

diff --git a/src/main/java/org/javacs/JavaCompilerService.java b/src/main/java/org/javacs/JavaCompilerService.java
index b4838da..f499a98 100644
--- a/src/main/java/org/javacs/JavaCompilerService.java
+++ b/src/main/java/org/javacs/JavaCompilerService.java
@@ -1160,8 +1160,38 @@ public class JavaCompilerService {
         return Optional.empty();
     }
 
+    private static Optional<Ref> ref(TreePath from, JavacTask task) {
+        var trees = Trees.instance(task);
+        var pos = trees.getSourcePositions();
+        var root = from.getCompilationUnit();
+        var lines = root.getLineMap();
+        var to = trees.getElement(from);
+        // Skip elements we can't find
+        if (to == null) {
+            LOG.warning(String.format("No element for %s", from.getLeaf()));
+            return Optional.empty();
+        }
+        // Skip non-methods
+        if (!(to instanceof ExecutableElement || to instanceof TypeElement || to instanceof VariableElement)) {
+            return Optional.empty();
+        }
+        // TODO skip anything not on source path
+        var toEl = idEl(to);
+        long start = pos.getStartPosition(root, from.getLeaf()), end = pos.getEndPosition(root, from.getLeaf());
+        if (start == -1) {
+            LOG.warning(String.format("No position for %s", from.getLeaf()));
+            return Optional.empty();
+        }
+        if (end == -1) end = start + from.getLeaf().toString().length();
+        int startLine = (int) lines.getLineNumber(start), startCol = (int) lines.getColumnNumber(start);
+        int endLine = (int) lines.getLineNumber(end), endCol = (int) lines.getColumnNumber(end);
+        var fromFile = root.getSourceFile().toUri();
+        var ref = new Ref(fromFile, startLine, startCol, endLine, endCol, toEl);
+        return Optional.of(ref);
+    }
+
     /** Compile `file` and locate `e` in it */
-    private Optional<TreePath> findIn(Element e, Path file, String contents) {
+    private Optional<Ref> findIn(Element e, Path file, String contents) {
         var task = singleFileTask(file.toUri(), contents);
         CompilationUnitTree tree;
         try {
@@ -1172,7 +1202,7 @@ public class JavaCompilerService {
         }
         var trees = Trees.instance(task);
         class Find extends TreePathScanner<Void, Void> {
-            Optional<TreePath> found = Optional.empty();
+            Optional<Ref> found = Optional.empty();
 
             boolean toStringEquals(Object left, Object right) {
                 return Objects.equals(Objects.toString(left, ""), Objects.toString(right, ""));
@@ -1188,7 +1218,7 @@ public class JavaCompilerService {
 
             void check() {
                 if (sameSymbol()) {
-                    found = Optional.of(getCurrentPath());
+                    found = ref(getCurrentPath(), cache.task);
                 }
             }
 
@@ -1210,7 +1240,7 @@ public class JavaCompilerService {
                 return super.visitVariable(node, aVoid);
             }
 
-            Optional<TreePath> run() {
+            Optional<Ref> run() {
                 scan(tree, null);
                 return found;
             }
@@ -1218,7 +1248,7 @@ public class JavaCompilerService {
         return new Find().run();
     }
 
-    public Optional<TreePath> definition(URI file, int line, int character, Function<URI, String> contents) {
+    public Optional<Ref> definition(URI file, int line, int character, Function<URI, String> contents) {
         recompile(file, contents.apply(file), line, character);
 
         var trees = Trees.instance(cache.task);
@@ -1455,7 +1485,7 @@ public class JavaCompilerService {
                 void check(TreePath from) {
                     var found = trees.getElement(from);
                     if (sameSymbol(found, to)) {
-                        ref(from).ifPresent(results::add);
+                        ref(from, task).ifPresent(results::add);
                     }
                 }
 
@@ -1481,41 +1511,11 @@ public class JavaCompilerService {
             return results;
         }
 
-        private Optional<Ref> ref(TreePath from) {
-            var trees = Trees.instance(task);
-            var pos = trees.getSourcePositions();
-            var root = from.getCompilationUnit();
-            var lines = root.getLineMap();
-            var to = trees.getElement(from);
-            // Skip elements we can't find
-            if (to == null) {
-                LOG.warning(String.format("No element for %s", from.getLeaf()));
-                return Optional.empty();
-            }
-            // Skip non-methods
-            if (!(to instanceof ExecutableElement)) {
-                return Optional.empty();
-            }
-            // TODO skip anything not on source path
-            var toEl = idEl(to);
-            long start = pos.getStartPosition(root, from.getLeaf()), end = pos.getEndPosition(root, from.getLeaf());
-            if (start == -1) {
-                LOG.warning(String.format("No position for %s", from.getLeaf()));
-                return Optional.empty();
-            }
-            if (end == -1) end = start + from.getLeaf().toString().length();
-            int startLine = (int) lines.getLineNumber(start), startCol = (int) lines.getColumnNumber(start);
-            int endLine = (int) lines.getLineNumber(end), endCol = (int) lines.getColumnNumber(end);
-            var fromFile = root.getSourceFile().toUri();
-            var ref = new Ref(fromFile, startLine, startCol, endLine, endCol, toEl);
-            return Optional.of(ref);
-        }
-
         List<Ref> index(CompilationUnitTree root) {
             var refs = new ArrayList<Ref>();
             class IndexFile extends TreePathScanner<Void, Void> {
                 void check(TreePath from) {
-                    ref(from).ifPresent(refs::add);
+                    ref(from, task).ifPresent(refs::add);
                 }
 
                 @Override
diff --git a/src/main/java/org/javacs/JavaTextDocumentService.java b/src/main/java/org/javacs/JavaTextDocumentService.java
index 0bfe4e5..88325d1 100644
--- a/src/main/java/org/javacs/JavaTextDocumentService.java
+++ b/src/main/java/org/javacs/JavaTextDocumentService.java
@@ -5,7 +5,6 @@ import com.sun.source.doctree.DocCommentTree;
 import com.sun.source.doctree.DocTree;
 import com.sun.source.doctree.ParamTree;
 import com.sun.source.tree.MethodTree;
-import com.sun.source.util.TreePath;
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.StringReader;
@@ -344,30 +343,6 @@ class JavaTextDocumentService implements TextDocumentService {
         return CompletableFuture.completedFuture(help);
     }
 
-    // TODO this is very duplicative with code in JavaCompilerService
-    // Maybe we shouldn't be passing around TreePath
-    private Optional<Location> location(TreePath p) {
-        var trees = server.compiler.trees();
-        var pos = trees.getSourcePositions();
-        var cu = p.getCompilationUnit();
-        var lines = cu.getLineMap();
-        long start = pos.getStartPosition(cu, p.getLeaf()), end = pos.getEndPosition(cu, p.getLeaf());
-        if (start == -1) {
-            LOG.warning(String.format("No position for %s", p.getLeaf()));
-            return Optional.empty();
-        }
-        if (end == -1) {
-            end = start + p.getLeaf().toString().length();
-        }
-        int startLine = (int) lines.getLineNumber(start) - 1, startCol = (int) lines.getColumnNumber(start) - 1;
-        int endLine = (int) lines.getLineNumber(end) - 1, endCol = (int) lines.getColumnNumber(end) - 1;
-        var dUri = cu.getSourceFile().toUri();
-        var startPos = new Position(startLine, startCol);
-        var endPos = new Position(endLine, endCol);
-        var loc = new Location(dUri.toString(), new Range(startPos, endPos));
-        return Optional.of(loc);
-    }
-
     @Override
     public CompletableFuture<List<? extends Location>> definition(TextDocumentPositionParams position) {
         var uri = URI.create(position.getTextDocument().getUri());
@@ -376,9 +351,8 @@ class JavaTextDocumentService implements TextDocumentService {
         Function<URI, String> findContent = f -> contents(f).content;
         var def = server.compiler.definition(uri, line, column, findContent);
         if (!def.isPresent()) return CompletableFuture.completedFuture(List.of());
-        var loc = location(def.get());
-        if (!loc.isPresent()) return CompletableFuture.completedFuture(List.of());
-        return CompletableFuture.completedFuture(List.of(loc.get()));
+        var loc = asLocation(def.get());
+        return CompletableFuture.completedFuture(List.of(loc));
     }
 
     class ReportProgress implements ReportReferencesProgress, AutoCloseable {
@@ -493,7 +467,11 @@ class JavaTextDocumentService implements TextDocumentService {
             var pos = maybePos.get();
             var range = asRange(pos);
             var message = pos.memberName.isPresent() ? "Run Test" : "Run All Tests";
-            var command = new Command(message, "java.command.test.run", List.of(uri, pos.className, pos.memberName));
+            var command =
+                    new Command(
+                            message,
+                            "java.command.test.run",
+                            Arrays.asList(uri, pos.className.orElse(null), pos.memberName.orElse(null)));
             var lens = new CodeLens(range, command, null);
             result.add(lens);
         }
diff --git a/src/test/java/org/javacs/GotoTest.java b/src/test/java/org/javacs/GotoTest.java
index 6c64834..e52825d 100644
--- a/src/test/java/org/javacs/GotoTest.java
+++ b/src/test/java/org/javacs/GotoTest.java
@@ -4,122 +4,118 @@ import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.assertThat;
 
 import java.net.URI;
+import java.nio.file.Paths;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
-import java.util.logging.Logger;
 import org.eclipse.lsp4j.Location;
 import org.eclipse.lsp4j.Position;
-import org.eclipse.lsp4j.Range;
 import org.eclipse.lsp4j.TextDocumentIdentifier;
 import org.eclipse.lsp4j.TextDocumentPositionParams;
 import org.junit.Ignore;
 import org.junit.Test;
 
 public class GotoTest {
-    private static final Logger LOG = Logger.getLogger("main");
     private static final String file = "/org/javacs/example/Goto.java";
-    private static final URI uri = FindResource.uri(file),
-            other = FindResource.uri("/org/javacs/example/GotoOther.java");
     private static final String defaultConstructorFile = "/org/javacs/example/GotoDefaultConstructor.java";
-    private static final URI defaultConstructorUri = FindResource.uri(defaultConstructorFile);
 
     @Test
     public void localVariable() {
-        List<? extends Location> suggestions = doGoto(file, 9, 8);
+        var suggestions = doGoto(file, 9, 8);
 
-        assertThat(suggestions, contains(location(uri, 4, 8, 4, 21)));
+        assertThat(suggestions, contains("Goto.java:5(9)"));
     }
 
     @Test
     public void defaultConstructor() {
-        List<? extends Location> suggestions = doGoto(defaultConstructorFile, 4, 45);
+        var suggestions = doGoto(defaultConstructorFile, 4, 45);
 
-        assertThat(suggestions, contains(location(defaultConstructorUri, 2, 0, 6, 1)));
+        assertThat(suggestions, contains("GotoDefaultConstructor.java:3(1)"));
     }
 
     @Test
     public void constructor() {
-        List<? extends Location> suggestions = doGoto(file, 10, 20);
+        var suggestions = doGoto(file, 10, 20);
 
-        assertThat(suggestions, contains(location(uri, 2, 0, 47, 1)));
+        assertThat(suggestions, contains("Goto.java:3(1)"));
     }
 
     @Test
     public void className() {
-        List<? extends Location> suggestions = doGoto(file, 15, 8);
+        var suggestions = doGoto(file, 15, 8);
 
-        assertThat(suggestions, contains(location(uri, 2, 0, 47, 1)));
+        assertThat(suggestions, contains("Goto.java:3(1)"));
     }
 
     @Test
     public void staticField() {
-        List<? extends Location> suggestions = doGoto(file, 12, 21);
+        var suggestions = doGoto(file, 12, 21);
 
-        assertThat(suggestions, contains(location(uri, 35, 4, 35, 37)));
+        assertThat(suggestions, contains("Goto.java:36(5)"));
     }
 
     @Test
     public void field() {
-        List<? extends Location> suggestions = doGoto(file, 13, 21);
+        var suggestions = doGoto(file, 13, 21);
 
-        assertThat(suggestions, contains(location(uri, 36, 4, 36, 24)));
+        assertThat(suggestions, contains("Goto.java:37(5)"));
     }
 
     @Test
     public void staticMethod() {
-        List<? extends Location> suggestions = doGoto(file, 15, 13);
+        var suggestions = doGoto(file, 15, 13);
 
-        assertThat(suggestions, contains(location(uri, 37, 4, 39, 5)));
+        assertThat(suggestions, contains("Goto.java:38(5)"));
     }
 
     @Test
     public void method() {
-        List<? extends Location> suggestions = doGoto(file, 16, 13);
+        var suggestions = doGoto(file, 16, 13);
 
-        assertThat(suggestions, contains(location(uri, 40, 4, 42, 5)));
+        assertThat(suggestions, contains("Goto.java:41(5)"));
     }
 
     @Test
     public void staticMethodReference() {
-        List<? extends Location> suggestions = doGoto(file, 18, 26);
+        var suggestions = doGoto(file, 18, 26);
 
-        assertThat(suggestions, contains(location(uri, 37, 4, 39, 5)));
+        assertThat(suggestions, contains("Goto.java:38(5)"));
     }
 
     @Test
     public void methodReference() {
-        List<? extends Location> suggestions = doGoto(file, 19, 26);
+        var suggestions = doGoto(file, 19, 26);
 
-        assertThat(suggestions, contains(location(uri, 40, 4, 42, 5)));
+        assertThat(suggestions, contains("Goto.java:41(5)"));
     }
 
     @Test
     public void otherStaticMethod() {
-        List<? extends Location> suggestions = doGoto(file, 28, 24);
+        var suggestions = doGoto(file, 28, 24);
 
-        assertThat(suggestions, contains(hasProperty("uri", equalTo(other.toString()))));
+        assertThat(suggestions, contains(startsWith("GotoOther.java:")));
     }
 
     @Test
     public void otherMethod() {
-        List<? extends Location> suggestions = doGoto(file, 29, 17);
+        var suggestions = doGoto(file, 29, 17);
 
-        assertThat(suggestions, contains(hasProperty("uri", equalTo(other.toString()))));
+        assertThat(suggestions, contains(startsWith("GotoOther.java:")));
     }
 
     @Test
     public void otherCompiledFile() {
-        List<? extends Location> suggestions = doGoto(file, 28, 24);
+        var suggestions = doGoto(file, 28, 24);
 
-        assertThat(suggestions, contains(hasProperty("uri", equalTo(other.toString()))));
+        assertThat(suggestions, contains(startsWith("GotoOther.java:")));
     }
 
     @Test
     @Ignore // TODO
     public void typeParam() {
-        List<? extends Location> suggestions = doGoto(file, 45, 11);
+        var suggestions = doGoto(file, 45, 11);
 
-        assertThat(suggestions, contains(location(uri, 2, 18, 2, 23)));
+        assertThat(suggestions, contains("Goto.java:3(19)"));
     }
 
     @Test
@@ -130,29 +126,9 @@ public class GotoTest {
         assertThat(doGoto(file, 5, 35), not(empty()));
     }
 
-    private Location location(URI uri, int startRow, int startColumn, int endRow, int endColumn) {
-        Position start = new Position();
-        start.setLine(startRow);
-        start.setCharacter(startColumn);
-
-        Position end = new Position();
-        end.setLine(endRow);
-        end.setCharacter(endColumn);
-
-        Range range = new Range();
-        range.setStart(start);
-        range.setEnd(end);
-
-        Location location = new Location();
-        location.setUri(uri.toString());
-        location.setRange(range);
-
-        return location;
-    }
-
     private static final JavaLanguageServer server = LanguageServerFixture.getJavaLanguageServer();
 
-    private List<? extends Location> doGoto(String file, int row, int column) {
+    private List<String> doGoto(String file, int row, int column) {
         TextDocumentIdentifier document = new TextDocumentIdentifier();
 
         document.setUri(FindResource.uri(file).toString());
@@ -167,10 +143,19 @@ public class GotoTest {
         p.setTextDocument(document);
         p.setPosition(position);
 
+        // TODO extends is not coloring correctly
+        List<? extends Location> locations;
         try {
-            return server.getTextDocumentService().definition(p).get();
+            locations = server.getTextDocumentService().definition(p).get();
         } catch (InterruptedException | ExecutionException e) {
             throw new RuntimeException(e);
         }
+        var strings = new ArrayList<String>();
+        for (var l : locations) {
+            var fileName = Paths.get(URI.create(l.getUri())).getFileName();
+            var start = l.getRange().getStart();
+            strings.add(String.format("%s:%d(%d)", fileName, start.getLine() + 1, start.getCharacter() + 1));
+        }
+        return strings;
     }
 }