Prune reference and definition searches

12 files changed, 203 insertions, 31 deletions

diff --git a/scripts/benchmark.sh b/scripts/benchmark.sh
index b816e0c..bc896d1 100755
--- a/scripts/benchmark.sh
+++ b/scripts/benchmark.sh
@@ -11,7 +11,7 @@ mvn test-compile
 mvn dependency:build-classpath -DincludeScope=test -Dmdep.outputFile=scripts/classpath.txt
 
 # Run the benchmark
-java -cp $(cat scripts/classpath.txt):target/classes:target/test-classes --illegal-access=warn org.openjdk.jmh.Main
+java -cp $(cat scripts/classpath.txt):target/classes:target/test-classes --illegal-access=warn org.openjdk.jmh.Main BenchmarkPruner
 
 # Clean up
 rm scripts/classpath.txt
\ No newline at end of file
diff --git a/src/main/java/org/javacs/CompileBatch.java b/src/main/java/org/javacs/CompileBatch.java
index 96ab36e..65e3145 100644
--- a/src/main/java/org/javacs/CompileBatch.java
+++ b/src/main/java/org/javacs/CompileBatch.java
@@ -96,7 +96,12 @@ public class CompileBatch {
                 return Optional.ofNullable(el);
             }
         }
-        throw new RuntimeException("File " + uri + " isn't in batch " + roots);
+        // Somehow, uri was not in batch
+        var names = new StringJoiner(", ");
+        for (var r : roots) {
+            names.add(Parser.fileName(r.getSourceFile().toUri()));
+        }
+        throw new RuntimeException("File " + uri + " isn't in batch " + names);
     }
 
     public Optional<List<TreePath>> definitions(Element el) {
diff --git a/src/main/java/org/javacs/JavaLanguageServer.java b/src/main/java/org/javacs/JavaLanguageServer.java
index 4932cc1..41db831 100644
--- a/src/main/java/org/javacs/JavaLanguageServer.java
+++ b/src/main/java/org/javacs/JavaLanguageServer.java
@@ -775,14 +775,15 @@ class JavaLanguageServer extends LanguageServer {
         // Compile all files that *might* contain definitions of fromEl
         var toFiles = compiler.potentialDefinitions(toEl.get());
         if (toFiles.isEmpty()) return Optional.of(List.of());
-        var batch = compiler.compileBatch(latestText(toFiles));
+        var name = toEl.get().getSimpleName().toString();
+        var batch = compiler.compileBatch(pruneWord(toFiles, name));
 
         // Find fromEl again, so that we have an Element from the current batch
         var fromElAgain = batch.element(fromUri, fromLine, fromColumn).get();
 
         // Find all definitions of fromElAgain
         var toTreePaths = batch.definitions(fromElAgain);
-        if (toTreePaths.isEmpty()) return Optional.empty();
+        if (!toTreePaths.isPresent()) return Optional.empty();
         var result = new ArrayList<Location>();
         for (var path : toTreePaths.get()) {
             var toUri = path.getCompilationUnit().getSourceFile().toUri();
@@ -845,7 +846,8 @@ class JavaLanguageServer extends LanguageServer {
         // Compile all files that *might* contain references to toEl
         var fromFiles = compiler.potentialReferences(toEl.get());
         if (fromFiles.isEmpty()) return Optional.of(List.of());
-        var batch = compiler.compileBatch(latestText(fromFiles));
+        var name = toEl.get().getSimpleName().toString();
+        var batch = compiler.compileBatch(pruneWord(fromFiles, name));
 
         // Find toEl again, so that we have an Element from the current batch
         var toElAgain = batch.element(toUri, toLine, toColumn).get();
@@ -853,7 +855,7 @@ class JavaLanguageServer extends LanguageServer {
         // Find all references to toElAgain
         // TODO this should references to supers of toEl
         var fromTreePaths = batch.references(toElAgain);
-        if (fromTreePaths.isEmpty()) return Optional.empty();
+        if (!fromTreePaths.isPresent()) return Optional.empty();
         var result = new ArrayList<Location>();
         for (var path : fromTreePaths.get()) {
             var fromUri = path.getCompilationUnit().getSourceFile().toUri();
@@ -868,6 +870,16 @@ class JavaLanguageServer extends LanguageServer {
         return Optional.of(result);
     }
 
+    private List<JavaFileObject> pruneWord(List<URI> files, String word) {
+        var sources = new ArrayList<JavaFileObject>();
+        for (var f : files) {
+            var contents = contents(f).content;
+            var pruned = Pruner.prune(f, contents, word);
+            sources.add(new StringFileObject(pruned, f));
+        }
+        return sources;
+    }
+
     private List<JavaFileObject> latestText(List<URI> files) {
         var sources = new ArrayList<JavaFileObject>();
         for (var f : files) {
diff --git a/src/main/java/org/javacs/Profiler.java b/src/main/java/org/javacs/Profiler.java
index 7285732..742ee35 100644
--- a/src/main/java/org/javacs/Profiler.java
+++ b/src/main/java/org/javacs/Profiler.java
@@ -16,6 +16,7 @@ class Profiler implements TaskListener {
     public void started(TaskEvent e) {
         started.put(e.getKind(), Instant.now());
         files.add(e.getSourceFile().toUri());
+        // TODO log file name when we compile something that wasn't in the batch
     }
 
     @Override
diff --git a/src/main/java/org/javacs/Pruner.java b/src/main/java/org/javacs/Pruner.java
index 08f2553..bf7203b 100644
--- a/src/main/java/org/javacs/Pruner.java
+++ b/src/main/java/org/javacs/Pruner.java
@@ -1,32 +1,45 @@
 package org.javacs;
 
 import com.sun.source.tree.*;
+import com.sun.source.util.SourcePositions;
 import com.sun.source.util.TreeScanner;
 import com.sun.source.util.Trees;
 import java.io.IOException;
 import java.net.URI;
+import java.util.ArrayList;
+import java.util.regex.Pattern;
 
 class Pruner {
-    static String prune(URI file, String contents, int line, int character) {
-        var task = Parser.parseTask(new StringFileObject(contents, file));
-        CompilationUnitTree root;
-        try {
-            root = task.parse().iterator().next();
-        } catch (IOException e) {
-            throw new RuntimeException(e);
-        }
-        var buffer = new StringBuilder(contents);
-        var sourcePositions = Trees.instance(task).getSourcePositions();
-        var lines = root.getLineMap();
-        var cursor = lines.getPosition(line, character);
+    private static String prune(CompilationUnitTree root, SourcePositions pos, StringBuilder buffer, long[] offsets) {
 
         class Scan extends TreeScanner<Void, Void> {
             boolean erasedAfterCursor = false;
 
             boolean containsCursor(Tree node) {
-                long start = sourcePositions.getStartPosition(root, node),
-                        end = sourcePositions.getEndPosition(root, node);
-                return start <= cursor && cursor <= end;
+                var start = pos.getStartPosition(root, node);
+                var end = pos.getEndPosition(root, node);
+                for (var cursor : offsets) {
+                    if (start <= cursor && cursor <= end) {
+                        return true;
+                    }
+                }
+                return false;
+            }
+
+            long lastCursorIn(Tree node) {
+                var start = pos.getStartPosition(root, node);
+                var end = pos.getEndPosition(root, node);
+                long last = -1;
+                for (var cursor : offsets) {
+                    if (start <= cursor && cursor <= end) {
+                        last = cursor;
+                    }
+                }
+                if (last == -1) {
+                    throw new RuntimeException(
+                            String.format("No cursor in %s is between %d and %d", offsets, start, end));
+                }
+                return last;
             }
 
             void erase(long start, long end) {
@@ -42,26 +55,27 @@ class Pruner {
             }
 
             @Override
-            public Void visitImport(ImportTree node, Void aVoid) {
+            public Void visitImport(ImportTree node, Void __) {
                 // Erase 'static' keyword so autocomplete works better
                 if (containsCursor(node) && node.isStatic()) {
-                    var start = (int) sourcePositions.getStartPosition(root, node);
+                    var start = (int) pos.getStartPosition(root, node);
                     start = buffer.indexOf("static", start);
                     var end = start + "static".length();
                     erase(start, end);
                 }
 
-                return super.visitImport(node, aVoid);
+                return super.visitImport(node, null);
             }
 
             @Override
-            public Void visitBlock(BlockTree node, Void aVoid) {
+            public Void visitBlock(BlockTree node, Void __) {
                 if (containsCursor(node)) {
-                    super.visitBlock(node, aVoid);
+                    super.visitBlock(node, null);
                     // When we find the deepest block that includes the cursor
                     if (!erasedAfterCursor) {
+                        var cursor = lastCursorIn(node);
                         var start = cursor;
-                        var end = sourcePositions.getEndPosition(root, node);
+                        var end = pos.getEndPosition(root, node);
                         if (end >= buffer.length()) end = buffer.length() - 1;
                         // Find the next line
                         while (start < end && buffer.charAt((int) start) != '\n') start++;
@@ -74,8 +88,8 @@ class Pruner {
                 } else if (!node.getStatements().isEmpty()) {
                     var first = node.getStatements().get(0);
                     var last = node.getStatements().get(node.getStatements().size() - 1);
-                    var start = sourcePositions.getStartPosition(root, first);
-                    var end = sourcePositions.getEndPosition(root, last);
+                    var start = pos.getStartPosition(root, first);
+                    var end = pos.getEndPosition(root, last);
                     if (end >= buffer.length()) end = buffer.length() - 1;
                     erase(start, end);
                 }
@@ -92,4 +106,47 @@ class Pruner {
 
         return buffer.toString();
     }
+
+    static String prune(URI file, String contents, int line, int character) {
+        // Parse file
+        var task = Parser.parseTask(new StringFileObject(contents, file));
+        CompilationUnitTree root;
+        try {
+            root = task.parse().iterator().next();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+        // Erase all blocks that don't include line:character
+        var lines = root.getLineMap();
+        var cursor = lines.getPosition(line, character);
+        var pos = Trees.instance(task).getSourcePositions();
+        var buffer = new StringBuilder(contents);
+        return prune(root, pos, buffer, new long[] {cursor});
+    }
+
+    static String prune(URI file, String contents, String name) {
+        // Find all occurrences of name in contents
+        var list = new ArrayList<Long>();
+        var pattern = Pattern.compile("\\b" + Pattern.quote(name) + "\\b");
+        var matcher = pattern.matcher(contents);
+        while (matcher.find()) {
+            list.add((long) matcher.start());
+        }
+        var offsets = new long[list.size()];
+        for (var i = 0; i < list.size(); i++) {
+            offsets[i] = list.get(i);
+        }
+        // Parse file
+        var task = Parser.parseTask(new StringFileObject(contents, file));
+        CompilationUnitTree root;
+        try {
+            root = task.parse().iterator().next();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+        // Erase all blocks that don't contain name
+        var buffer = new StringBuilder(contents);
+        var pos = Trees.instance(task).getSourcePositions();
+        return prune(root, pos, buffer, offsets);
+    }
 }
diff --git a/src/main/java/org/javacs/lsp/LSP.java b/src/main/java/org/javacs/lsp/LSP.java
index 55f1eca..4663e0c 100644
--- a/src/main/java/org/javacs/lsp/LSP.java
+++ b/src/main/java/org/javacs/lsp/LSP.java
@@ -415,6 +415,7 @@ public class LSP {
                 }
             } catch (Exception e) {
                 LOG.log(Level.SEVERE, e.getMessage(), e);
+                // TODO send failure response if r was a request
             }
         }
     }
diff --git a/src/test/java/org/javacs/ErrorProneBenchmark.java b/src/test/java/org/javacs/BenchmarkErrorProne.java
index 6e0b7ba..ff3f64e 100644
--- a/src/test/java/org/javacs/ErrorProneBenchmark.java
+++ b/src/test/java/org/javacs/BenchmarkErrorProne.java
@@ -20,7 +20,7 @@ import org.openjdk.jmh.annotations.*;
 @Warmup(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
 @Measurement(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
 @Fork(1)
-public class ErrorProneBenchmark {
+public class BenchmarkErrorProne {
     private static Path file = Paths.get(FindResource.uri("/org/javacs/example/BenchmarkStringSearch.java"));
 
     @State(Scope.Thread)
diff --git a/src/test/java/org/javacs/BenchmarkPruner.java b/src/test/java/org/javacs/BenchmarkPruner.java
new file mode 100644
index 0000000..c212cd8
--- /dev/null
+++ b/src/test/java/org/javacs/BenchmarkPruner.java
@@ -0,0 +1,63 @@
+package org.javacs;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.openjdk.jmh.annotations.*;
+
+@Warmup(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
+@Fork(1)
+public class BenchmarkPruner {
+    private static Path sourceRoot = Paths.get("src/main/java").normalize();
+    private static List<StringFileObject> files = files(false);
+    private static List<StringFileObject> pruned = files(true);
+
+    private static List<StringFileObject> files(boolean prune) {
+        try {
+            var files = new ArrayList<StringFileObject>();
+            var dir = Paths.get("src/main/java/org/javacs").normalize();
+            var it = Files.list(dir).iterator();
+            while (it.hasNext()) {
+                var file = it.next();
+                if (!Files.isRegularFile(file)) continue;
+                var contents = String.join("\n", Files.readAllLines(file));
+                ;
+                if (prune) {
+                    contents = Pruner.prune(file.toUri(), contents, "isWord");
+                }
+                files.add(new StringFileObject(contents, file.toUri()));
+            }
+            return files;
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private static JavaCompilerService createCompiler() {
+        return new JavaCompilerService(
+                Collections.singleton(sourceRoot),
+                () -> {
+                    throw new RuntimeException("Unimplemented");
+                },
+                Collections.emptySet(),
+                Collections.emptySet());
+    }
+
+    @Benchmark
+    public void pruned() {
+        var compiler = createCompiler();
+        compiler.compileBatch(pruned);
+    }
+
+    @Benchmark
+    public void plain() {
+        var compiler = createCompiler();
+        compiler.compileBatch(files);
+    }
+}
diff --git a/src/test/java/org/javacs/StringSearchBenchmark.java b/src/test/java/org/javacs/BenchmarkStringSearch.java
index e044812..1c453d2 100644
--- a/src/test/java/org/javacs/StringSearchBenchmark.java
+++ b/src/test/java/org/javacs/BenchmarkStringSearch.java
@@ -9,7 +9,7 @@ import org.openjdk.jmh.annotations.*;
 @Warmup(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
 @Measurement(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
 @Fork(1)
-public class StringSearchBenchmark {
+public class BenchmarkStringSearch {
     private static final Path largeFile = Paths.get(FindResource.uri("/org/javacs/example/LargeFile.java")),
             smallFile = Paths.get(FindResource.uri("/org/javacs/example/Goto.java"));
     // "removeMethodBodies" appears late in the file, so stopping early will not be very effective
diff --git a/src/test/java/org/javacs/PrunerTest.java b/src/test/java/org/javacs/PrunerTest.java
index 24098be..50e511f 100644
--- a/src/test/java/org/javacs/PrunerTest.java
+++ b/src/test/java/org/javacs/PrunerTest.java
@@ -36,4 +36,11 @@ public class PrunerTest {
         var expected = contents("PruneDot_erased.java");
         assertThat(actual, equalToIgnoringWhiteSpace(expected));
     }
+
+    @Test
+    public void pruneWords() {
+        var actual = Pruner.prune(URI.create("/PruneWords.java"), contents("PruneWords.java"), "word");
+        var expected = contents("PruneWords_erased.java");
+        assertThat(actual, equalToIgnoringWhiteSpace(expected));
+    }
 }
diff --git a/src/test/test-project/simple/PruneWords.java b/src/test/test-project/simple/PruneWords.java
new file mode 100644
index 0000000..3d0e2bd
--- /dev/null
+++ b/src/test/test-project/simple/PruneWords.java
@@ -0,0 +1,13 @@
+class PruneWords {
+    void keepThis() {
+        word();
+    }
+
+    void eraseThis() {
+        notword();
+    }
+
+    void keepThisToo() {
+        word();
+    }
+}
\ No newline at end of file
diff --git a/src/test/test-project/simple/PruneWords_erased.java b/src/test/test-project/simple/PruneWords_erased.java
new file mode 100644
index 0000000..873f9cf
--- /dev/null
+++ b/src/test/test-project/simple/PruneWords_erased.java
@@ -0,0 +1,13 @@
+class PruneWords {
+    void keepThis() {
+        word();
+    }
+
+    void eraseThis() {
+
+    }
+
+    void keepThisToo() {
+        word();
+    }
+}
\ No newline at end of file