Re-use work across code lenses

4 files changed, 174 insertions, 35 deletions

diff --git a/src/main/java/org/javacs/CompileBatch.java b/src/main/java/org/javacs/CompileBatch.java
index abc6d53..d3524de 100644
--- a/src/main/java/org/javacs/CompileBatch.java
+++ b/src/main/java/org/javacs/CompileBatch.java
@@ -168,17 +168,34 @@ public class CompileBatch {
     }
 
     public Optional<List<TreePath>> references(Element to) {
-        LOG.info(String.format("Search for references to `%s` in %d files...", to, roots.size()));
+        var map = references(List.of(to));
+        if (map.size() > 1) {
+            throw new RuntimeException(String.format("Searched for `%s` but found multiple %s", to, map.keySet()));
+        }
+        // Return the only element in the map
+        for (var path : map.values()) {
+            return Optional.of(path);
+        }
+        // Map is empty, to must have been removed due to errors
+        return Optional.empty();
+    }
 
-        if (to.asType().getKind() == TypeKind.ERROR) {
-            LOG.info(String.format("...`%s` is an error type, giving up", to.asType()));
-            return Optional.empty();
+    public Map<Element, List<TreePath>> references(List<Element> toAny) {
+        LOG.info(String.format("Search for references to %d elements in %d files...", toAny.size(), roots.size()));
+
+        var els = new ArrayList<Element>();
+        for (var to : toAny) {
+            if (to.asType().getKind() == TypeKind.ERROR) {
+                LOG.info(String.format("...`%s` is an error type, giving up", to.asType()));
+                continue;
+            }
+            els.add(to);
         }
 
-        var refs = new ArrayList<TreePath>();
+        var refs = new HashMap<Element, List<TreePath>>();
         class FindReferences extends TreePathScanner<Void, Void> {
-            boolean sameSymbol(Element found) {
-                if (to.equals(found)) {
+            boolean sameSymbol(Element from, Element to) {
+                if (to.equals(from)) {
                     var uri = getCurrentPath().getCompilationUnit().getSourceFile().toUri();
                     var fileName = Parser.fileName(uri);
                     return true;
@@ -186,12 +203,12 @@ public class CompileBatch {
                 return false;
             }
 
-            boolean isSuperMethod(Element found) {
+            boolean isSuperMethod(Element from, Element to) {
                 if (!(to instanceof ExecutableElement)) return false;
-                if (!(found instanceof ExecutableElement)) return false;
+                if (!(from instanceof ExecutableElement)) return false;
                 var subMethod = (ExecutableElement) to;
                 var subType = (TypeElement) subMethod.getEnclosingElement();
-                var superMethod = (ExecutableElement) found;
+                var superMethod = (ExecutableElement) from;
                 // TODO need to check if class is compatible as well
                 if (elements.overrides(subMethod, superMethod, subType)) {
                     LOG.info(String.format("...`%s.%s` overrides `%s`", subType, subMethod, superMethod));
@@ -201,9 +218,14 @@ public class CompileBatch {
             }
 
             void check(TreePath from) {
-                var found = trees.getElement(from);
-                var match = sameSymbol(found) || isSuperMethod(found);
-                if (match) refs.add(from);
+                for (var to : els) {
+                    var fromEl = trees.getElement(from);
+                    var match = sameSymbol(fromEl, to) || isSuperMethod(fromEl, to);
+                    if (match) {
+                        var list = refs.computeIfAbsent(to, __ -> new ArrayList<>());
+                        list.add(from);
+                    }
+                }
             }
 
             @Override
@@ -234,7 +256,30 @@ public class CompileBatch {
         for (var r : roots) {
             finder.scan(r, null);
         }
-        return Optional.of(refs);
+        return refs;
+    }
+
+    /**
+     * Find all elements in `file` that get turned into code-lenses. This needs to match the result of
+     * `ParseFile#declarations`
+     */
+    public List<Element> declarations(URI file) {
+        for (var r : roots) {
+            if (!r.getSourceFile().toUri().equals(file)) continue;
+            var paths = ParseFile.declarations(r);
+            var els = new ArrayList<Element>();
+            for (var p : paths) {
+                var e = trees.getElement(p);
+                assert e != null;
+                els.add(e);
+            }
+            return els;
+        }
+        var message = new StringJoiner(", ");
+        for (var r : roots) {
+            message.add(Parser.fileName(r.getSourceFile().toUri()));
+        }
+        throw new RuntimeException(file + " is not in " + message);
     }
 
     public Optional<Range> range(TreePath path) {
diff --git a/src/main/java/org/javacs/JavaCompilerService.java b/src/main/java/org/javacs/JavaCompilerService.java
index 21adbc6..edf53cd 100644
--- a/src/main/java/org/javacs/JavaCompilerService.java
+++ b/src/main/java/org/javacs/JavaCompilerService.java
@@ -322,21 +322,24 @@ public class JavaCompilerService {
                     found.add(uri);
                 }
 
+                public boolean isName(Tree t) {
+                    if (t instanceof MemberSelectTree) {
+                        var select = (MemberSelectTree) t;
+                        return select.getIdentifier().contentEquals(findName);
+                    }
+                    if (t instanceof IdentifierTree) {
+                        var id = (IdentifierTree) t;
+                        return id.getName().contentEquals(findName);
+                    }
+                    return false;
+                }
+
                 @Override
                 public Void visitMethodInvocation(MethodInvocationTree t, Set<URI> found) {
                     // TODO try to disprove that this is a reference by looking at obvious special cases, like is the
                     // simple name of the type different?
                     var method = t.getMethodSelect();
-                    // outer.method()
-                    if (method instanceof MemberSelectTree) {
-                        var select = (MemberSelectTree) method;
-                        if (select.getIdentifier().contentEquals(findName)) add(found);
-                    }
-                    // method()
-                    if (method instanceof IdentifierTree) {
-                        var id = (IdentifierTree) method;
-                        if (id.getName().contentEquals(findName)) add(found);
-                    }
+                    if (isName(method)) add(found);
                     // Check other parts
                     return super.visitMethodInvocation(t, found);
                 }
@@ -346,6 +349,13 @@ public class JavaCompilerService {
                     if (t.getName().contentEquals(findName)) add(found);
                     return super.visitMemberReference(t, found);
                 }
+
+                @Override
+                public Void visitNewClass(NewClassTree t, Set<URI> found) {
+                    var cls = t.getIdentifier();
+                    if (isName(cls)) add(found);
+                    return super.visitNewClass(t, found);
+                }
             }
             return scanForPotentialReferences(to, new FindMethod());
         } else {
diff --git a/src/main/java/org/javacs/JavaLanguageServer.java b/src/main/java/org/javacs/JavaLanguageServer.java
index 8731fbb..2193788 100644
--- a/src/main/java/org/javacs/JavaLanguageServer.java
+++ b/src/main/java/org/javacs/JavaLanguageServer.java
@@ -348,6 +348,7 @@ class JavaLanguageServer extends LanguageServer {
         var column = position.position.character + 1;
         // Figure out what kind of completion we want to do
         var maybeCtx = compiler.parseFile(uri, content).completionContext(line, column);
+        // TODO don't complete inside of comments
         if (!maybeCtx.isPresent()) {
             var items = new ArrayList<CompletionItem>();
             for (var name : CompileFocus.TOP_LEVEL_KEYWORDS) {
@@ -848,7 +849,7 @@ class JavaLanguageServer extends LanguageServer {
         return sources;
     }
 
-    private List<JavaFileObject> latestText(List<URI> files) {
+    private List<JavaFileObject> latestText(Collection<URI> files) {
         var sources = new ArrayList<JavaFileObject>();
         for (var f : files) {
             sources.add(new StringFileObject(contents(f).content, f));
@@ -1039,22 +1040,96 @@ class JavaLanguageServer extends LanguageServer {
         // TODO if this gets too big, just show "Many references"
         var fromFiles = compiler.potentialReferences(toEl.get());
         fromFiles.add(toUri);
-        // TODO instead of pruning words (2x speedup at best),
-        // compile the whole file and save all references to all declarations in all open files.
-        // As soon as any open file is edited, discard this cache.
-        // This will re-use work across code lenses.
-        var batch = compiler.compileBatch(pruneWord(fromFiles, toEl.get()));
 
-        // Find toEl again, so that we have an Element from the current batch
-        var toElAgain = batch.element(toUri, toLine, toColumn).get();
+        // Make sure all fromFiles -> toUri references are in the cache
+        updateCountReferencesCache(toUri, fromFiles);
 
-        // Find all references to toElAgain
-        var fromTreePaths = batch.references(toElAgain);
-        var count = fromTreePaths.orElse(List.of()).size();
+        // Count up how many total references exist in fromFiles
+        var toPtr = new Ptr(toEl.get());
+        var count = 0;
+        for (var from : fromFiles) {
+            var cachedFileCounts = countReferencesCache.get(from);
+            count += cachedFileCounts.counts.getOrDefault(toPtr, 0);
+        }
         if (count == 1) return "1 reference";
         return String.format("%d references", count);
     }
 
+    /** countReferencesCache[file][ptr] is the number of references to ptr in file */
+    private Map<URI, CountReferences> countReferencesCache = new HashMap<>();
+
+    private static class CountReferences {
+        final Map<Ptr, Integer> counts = new HashMap<>();
+        final Instant created = Instant.now();
+    }
+
+    /** countReferencesCacheFile is the file pointed to by every ptr in countReferencesCache[_][ptr] */
+    private URI countReferencesCacheFile = URI.create("file:///NONE");
+
+    /** countReferencesCacheVersion is the version of countReferencesCacheFile that is currently cached */
+    private int countReferencesCacheVersion = -1;
+
+    private void updateCountReferencesCache(URI toFile, Collection<URI> fromFiles) {
+        // If cached file has changed, invalidate the whole cache
+        if (!toFile.equals(countReferencesCacheFile) || version(toFile) != countReferencesCacheVersion) {
+            LOG.info(String.format("Cache count-references %s", Parser.fileName(toFile)));
+            countReferencesCache.clear();
+            countReferencesCacheFile = toFile;
+            countReferencesCacheVersion = version(toFile);
+        }
+
+        // Figure out which from-files are out-of-date
+        var outOfDate = new HashSet<URI>();
+        for (var f : fromFiles) {
+            Instant modified;
+            try {
+                modified = Files.getLastModifiedTime(Paths.get(f)).toInstant();
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+            var expired =
+                    !countReferencesCache.containsKey(f) || countReferencesCache.get(f).created.isBefore(modified);
+            if (expired) {
+                countReferencesCache.remove(f);
+                outOfDate.add(f);
+            }
+        }
+
+        // Compile all out-of-date files
+        if (outOfDate.isEmpty()) return;
+        LOG.info(
+                String.format(
+                        "...%d files need to be re-counted for references to %s",
+                        outOfDate.size(), Parser.fileName(toFile)));
+        // TODO this extra file could be eliminated by remembering a List<Ptr> for the current file
+        outOfDate.add(toFile);
+        countReferencesCache.remove(toFile);
+        var batch = compiler.compileBatch(latestText(outOfDate));
+
+        // Find all declarations in toFile
+        var allEls = batch.declarations(toFile);
+
+        // Find all references to all declarations
+        var refs = batch.references(allEls);
+
+        // Update cached counts
+        for (var to : refs.keySet()) {
+            var toPtr = new Ptr(to);
+
+            for (var from : refs.get(to)) {
+                var fromUri = from.getCompilationUnit().getSourceFile().toUri();
+                var counts = countReferencesCache.computeIfAbsent(fromUri, __ -> new CountReferences());
+                var c = counts.counts.getOrDefault(toPtr, 0);
+                counts.counts.put(toPtr, c + 1);
+            }
+        }
+
+        // Ensure that all fromFiles are in the cache, even if they contain no references to toFile
+        for (var fromUri : fromFiles) {
+            countReferencesCache.computeIfAbsent(fromUri, __ -> new CountReferences());
+        }
+    }
+
     @Override
     public List<TextEdit> formatting(DocumentFormattingParams params) {
         updateHoverCache(params.textDocument.uri, contents(params.textDocument.uri).content);
@@ -1315,6 +1390,11 @@ class JavaLanguageServer extends LanguageServer {
         return activeDocuments.keySet();
     }
 
+    int version(URI openFile) {
+        if (!activeDocuments.containsKey(openFile)) return -1;
+        return activeDocuments.get(openFile).version;
+    }
+
     VersionedContent contents(URI openFile) {
         if (!isJavaFile(openFile)) {
             LOG.warning("Ignoring non-java file " + openFile);
diff --git a/src/main/java/org/javacs/ParseFile.java b/src/main/java/org/javacs/ParseFile.java
index 40ed816..78955ea 100644
--- a/src/main/java/org/javacs/ParseFile.java
+++ b/src/main/java/org/javacs/ParseFile.java
@@ -83,6 +83,10 @@ public class ParseFile {
     }
 
     public List<TreePath> declarations() {
+        return declarations(root);
+    }
+
+    static List<TreePath> declarations(CompilationUnitTree root) {
         var found = new ArrayList<TreePath>();
         class FindDeclarations extends TreePathScanner<Void, Void> {
             boolean isClass(Tree t) {