Reduce use of TreePath

3 files changed, 47 insertions, 40 deletions

diff --git a/src/main/java/org/javacs/JavaCompilerService.java b/src/main/java/org/javacs/JavaCompilerService.java
index 8a792fc..b4838da 100644
--- a/src/main/java/org/javacs/JavaCompilerService.java
+++ b/src/main/java/org/javacs/JavaCompilerService.java
@@ -1448,13 +1448,15 @@ public class JavaCompilerService {
                     && toStringEquals(to, from);
         }
 
-        List<TreePath> referencesToElement(CompilationUnitTree root, Element to) {
+        List<Ref> referencesToElement(CompilationUnitTree root, Element to) {
             var trees = Trees.instance(task);
-            var results = new ArrayList<TreePath>();
+            var results = new ArrayList<Ref>();
             class FindReferencesElement extends TreePathScanner<Void, Void> {
                 void check(TreePath from) {
                     var found = trees.getElement(from);
-                    if (sameSymbol(found, to)) results.add(from);
+                    if (sameSymbol(found, to)) {
+                        ref(from).ifPresent(results::add);
+                    }
                 }
 
                 @Override
@@ -1479,36 +1481,41 @@ public class JavaCompilerService {
             return results;
         }
 
-        List<Ref> index(CompilationUnitTree root) {
+        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) {
-                    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;
-                    }
-                    // Skip non-methods
-                    if (!(to instanceof ExecutableElement)) {
-                        return;
-                    }
-                    // 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;
-                    }
-                    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);
-                    refs.add(ref);
+                    ref(from).ifPresent(refs::add);
                 }
 
                 @Override
@@ -1578,7 +1585,7 @@ public class JavaCompilerService {
         return result;
     }
 
-    public List<TreePath> references(URI file, String contents, int line, int character, ReportReferencesProgress progress) {
+    public List<Ref> references(URI file, String contents, int line, int character, ReportReferencesProgress progress) {
         recompile(file, contents, -1, -1);
 
         var trees = Trees.instance(cache.task);
@@ -1597,7 +1604,7 @@ public class JavaCompilerService {
         progress.checkPotentialReferences(0, possible.size());
         // TODO optimize by pruning method bodies that don't contain potential references
         var batch = compileBatch(possible);
-        var result = new ArrayList<TreePath>();
+        var result = new ArrayList<Ref>();
         var nChecked = 0;
         for (var f : batch.roots) {
             result.addAll(batch.referencesToElement(f, toEl));
diff --git a/src/main/java/org/javacs/JavaTextDocumentService.java b/src/main/java/org/javacs/JavaTextDocumentService.java
index ef6635a..0bfe4e5 100644
--- a/src/main/java/org/javacs/JavaTextDocumentService.java
+++ b/src/main/java/org/javacs/JavaTextDocumentService.java
@@ -437,8 +437,7 @@ class JavaTextDocumentService implements TextDocumentService {
         Function<Integer, String> checkMessage = nPotential -> String.format("Check %,d files", nPotential);
         try (var progress = new ReportProgress(startMessage, scanMessage, checkMessage)) {
             for (var r : server.compiler.references(uri, content, line, column, progress)) {
-                var loc = location(r);
-                if (loc.isPresent()) result.add(loc.get());
+                result.add(asLocation(r));
             }
             return CompletableFuture.completedFuture(result);
         }
@@ -474,6 +473,13 @@ class JavaTextDocumentService implements TextDocumentService {
         return new Range(start, end);
     }
 
+    private Location asLocation(Ref ref) {
+        var start = new Position(ref.startLine - 1, ref.startCol - 1);
+        var end = new Position(ref.endLine - 1, ref.endCol - 1);
+        var range = new Range(start, end);
+        return new Location(ref.fromFile.toString(), range);
+    }
+
     private List<CodeLens> testMethods(URI uri) {
         var content = contents(uri).content;
         var tests = server.compiler.testMethods(uri, content);
diff --git a/src/test/java/org/javacs/JavaCompilerServiceTest.java b/src/test/java/org/javacs/JavaCompilerServiceTest.java
index 6023337..d683dba 100644
--- a/src/test/java/org/javacs/JavaCompilerServiceTest.java
+++ b/src/test/java/org/javacs/JavaCompilerServiceTest.java
@@ -225,15 +225,9 @@ public class JavaCompilerServiceTest {
         var file = "/GotoDefinition.java";
         var refs = compiler.references(resourceUri(file), contents(file), 6, 13, rrp);
         var stringify = new ArrayList<String>();
-        for (var t : refs) {
-            var unit = t.getCompilationUnit();
-            var name = Paths.get(unit.getSourceFile().toUri()).getFileName();
-            var trees = compiler.trees();
-            var pos = trees.getSourcePositions();
-            var lines = unit.getLineMap();
-            var start = pos.getStartPosition(unit, t.getLeaf());
-            var line = lines.getLineNumber(start);
-            stringify.add(String.format("%s:%d", name, line));
+        for (var r : refs) {
+            var fileName = Paths.get(r.fromFile).getFileName();
+            stringify.add(String.format("%s:%d", fileName, r.startLine));
         }
         assertThat(stringify, hasItem("GotoDefinition.java:3"));
         assertThat(stringify, not(hasItem("GotoDefinition.java:6")));