Re-use reference counts

9 files changed, 246 insertions, 158 deletions

diff --git a/TODOS.md b/TODOS.md
index c94e742..d6594e3 100644
--- a/TODOS.md
+++ b/TODOS.md
@@ -7,7 +7,4 @@
 - Go-to-subclasses
 - Test coverage codelens
 - Go-to-definition for overriding methods
-- Go-to-implementation for overridden methods
-
-## Polish
-- Show warning for unused local var, unused private method
\ No newline at end of file
+- Go-to-implementation for overridden methods
\ No newline at end of file
diff --git a/lib/extension.ts b/lib/extension.ts
index 2be8a9e..e739434 100644
--- a/lib/extension.ts
+++ b/lib/extension.ts
@@ -6,6 +6,7 @@ import * as FS from "fs";
 import { window, workspace, ExtensionContext, commands, tasks, Task, TaskExecution, ShellExecution, Uri, TaskDefinition, languages, IndentAction, Progress, ProgressLocation } from 'vscode';
 import {LanguageClient, LanguageClientOptions, ServerOptions, NotificationType} from "vscode-languageclient";
 
+// If we want to profile using VisualVM, we have to run the language server using regular java, not jlink
 const visualVm = false;
 
 /** Called when extension is activated */
diff --git a/src/main/java/org/javacs/Cache.java b/src/main/java/org/javacs/Cache.java
index b1c893d..6344044 100644
--- a/src/main/java/org/javacs/Cache.java
+++ b/src/main/java/org/javacs/Cache.java
@@ -70,4 +70,8 @@ class Cache<K, V> {
         }
         return map.get(k).value;
     }
+
+    void clear() {
+        map.clear();
+    }
 }
diff --git a/src/main/java/org/javacs/CompileBatch.java b/src/main/java/org/javacs/CompileBatch.java
index 6b94b73..d6b9bda 100644
--- a/src/main/java/org/javacs/CompileBatch.java
+++ b/src/main/java/org/javacs/CompileBatch.java
@@ -163,91 +163,32 @@ public class CompileBatch {
     }
 
     public Optional<List<TreePath>> references(Element to) {
-        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();
-    }
+        LOG.info(String.format("Search for references to %s...", to));
 
-    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);
+        // If to is an error, we won't be able to find anything
+        if (to.asType().getKind() == TypeKind.ERROR) {
+            LOG.info(String.format("...`%s` is an error type, giving up", to.asType()));
+            return Optional.empty();
         }
 
-        var refs = new HashMap<Element, List<TreePath>>();
-        class FindReferences extends TreePathScanner<Void, Void> {
-            boolean sameSymbol(Element from, Element to) {
-                return to.equals(from);
-            }
-
-            boolean isSuperMethod(Element from, Element to) {
-                if (!(to instanceof ExecutableElement)) return false;
-                if (!(from instanceof ExecutableElement)) return false;
-                var subMethod = (ExecutableElement) to;
-                var subType = (TypeElement) subMethod.getEnclosingElement();
-                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));
-                    return true;
-                }
-                return false;
-            }
-
-            void check(TreePath 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
-            public Void visitMemberReference(MemberReferenceTree t, Void __) {
-                check(getCurrentPath());
-                return super.visitMemberReference(t, null);
-            }
-
-            @Override
-            public Void visitMemberSelect(MemberSelectTree t, Void __) {
-                check(getCurrentPath());
-                return super.visitMemberSelect(t, null);
-            }
-
-            @Override
-            public Void visitIdentifier(IdentifierTree t, Void __) {
-                check(getCurrentPath());
-                return super.visitIdentifier(t, null);
-            }
-
-            @Override
-            public Void visitNewClass(NewClassTree t, Void __) {
-                check(getCurrentPath());
-                return super.visitNewClass(t, null);
-            }
+        // Otherwise, scan roots for references
+        List<TreePath> list = new ArrayList<TreePath>();
+        var map = Map.of(to, list);
+        var finder = new FindReferences(task);
+        for (var r : roots) {
+            finder.scan(r, map);
         }
-        var finder = new FindReferences();
+        LOG.info(String.format("...found %d references", list.size()));
+        return Optional.of(list);
+    }
+
+    public Index index(URI from, List<Element> toAny) {
         for (var r : roots) {
-            finder.scan(r, null);
+            if (r.getSourceFile().toUri().equals(from)) {
+                return new Index(task, r, parent.diags, toAny);
+            }
         }
-        LOG.info(String.format("...found %d references", refs.size()));
-        return refs;
+        throw new RuntimeException(from + " is not in compiled batch");
     }
 
     /**
diff --git a/src/main/java/org/javacs/CompileFile.java b/src/main/java/org/javacs/CompileFile.java
index 7051377..8bcef9c 100644
--- a/src/main/java/org/javacs/CompileFile.java
+++ b/src/main/java/org/javacs/CompileFile.java
@@ -47,6 +47,21 @@ public class CompileFile {
         return trees.getSourcePositions();
     }
 
+    /**
+     * Find all elements in `file` that get turned into code-lenses. This needs to match the result of
+     * `ParseFile#declarations`
+     */
+    public List<Element> declarations() {
+        var paths = ParseFile.declarations(root);
+        var els = new ArrayList<Element>();
+        for (var p : paths) {
+            var e = trees.getElement(p);
+            assert e != null;
+            els.add(e);
+        }
+        return els;
+    }
+
     public Optional<Element> element(int line, int character) {
         // LOG.info(String.format("Looking for element at %s(%d,%d)...", file.getPath(), line, character));
 
diff --git a/src/main/java/org/javacs/FindReferences.java b/src/main/java/org/javacs/FindReferences.java
new file mode 100644
index 0000000..f94240a
--- /dev/null
+++ b/src/main/java/org/javacs/FindReferences.java
@@ -0,0 +1,81 @@
+package org.javacs;
+
+import com.sun.source.tree.*;
+import com.sun.source.util.*;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.lang.model.element.*;
+import javax.lang.model.util.Elements;
+
+class FindReferences extends TreePathScanner<Void, Map<Element, List<TreePath>>> {
+    private final Trees trees;
+    private final Elements elements;
+
+    static Map<Element, List<TreePath>> empty(Collection<Element> to) {
+        var refs = new HashMap<Element, List<TreePath>>();
+        for (var el : to) {
+            refs.put(el, new ArrayList<>());
+        }
+        return refs;
+    }
+
+    FindReferences(JavacTask task) {
+        this.trees = Trees.instance(task);
+        this.elements = task.getElements();
+    }
+
+    boolean sameSymbol(Element from, Element to) {
+        return to.equals(from);
+    }
+
+    boolean isSuperMethod(Element from, Element to) {
+        if (!(to instanceof ExecutableElement)) return false;
+        if (!(from instanceof ExecutableElement)) return false;
+        var subMethod = (ExecutableElement) to;
+        var subType = (TypeElement) subMethod.getEnclosingElement();
+        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));
+            return true;
+        }
+        return false;
+    }
+
+    void check(TreePath from, Map<Element, List<TreePath>> refs) {
+        for (var to : refs.keySet()) {
+            var fromEl = trees.getElement(from);
+            var match = sameSymbol(fromEl, to) || isSuperMethod(fromEl, to);
+            if (match) {
+                refs.get(to).add(from);
+            }
+        }
+    }
+
+    @Override
+    public Void visitMemberReference(MemberReferenceTree t, Map<Element, List<TreePath>> refs) {
+        check(getCurrentPath(), refs);
+        return super.visitMemberReference(t, refs);
+    }
+
+    @Override
+    public Void visitMemberSelect(MemberSelectTree t, Map<Element, List<TreePath>> refs) {
+        check(getCurrentPath(), refs);
+        return super.visitMemberSelect(t, refs);
+    }
+
+    @Override
+    public Void visitIdentifier(IdentifierTree t, Map<Element, List<TreePath>> refs) {
+        check(getCurrentPath(), refs);
+        return super.visitIdentifier(t, refs);
+    }
+
+    @Override
+    public Void visitNewClass(NewClassTree t, Map<Element, List<TreePath>> refs) {
+        check(getCurrentPath(), refs);
+        return super.visitNewClass(t, refs);
+    }
+}
diff --git a/src/main/java/org/javacs/Index.java b/src/main/java/org/javacs/Index.java
new file mode 100644
index 0000000..6ab404d
--- /dev/null
+++ b/src/main/java/org/javacs/Index.java
@@ -0,0 +1,58 @@
+package org.javacs;
+
+import com.sun.source.tree.*;
+import com.sun.source.util.*;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import javax.lang.model.element.Element;
+import javax.tools.*;
+
+class Index {
+    /** map[ptr] is the number of references to ptr */
+    private final Map<Ptr, Integer> referenceCounts = new HashMap<>();
+    /** hasErrors is true if there were compilation errors when we created this index */
+    private boolean hasErrors;
+
+    Index(
+            JavacTask task,
+            CompilationUnitTree from,
+            Collection<Diagnostic<? extends JavaFileObject>> errors,
+            Collection<Element> to) {
+        // Scan from for references
+        var finder = new FindReferences(task);
+        var refs = new HashMap<Element, List<TreePath>>();
+        for (var el : to) {
+            refs.put(el, new ArrayList<>());
+        }
+        finder.scan(from, refs);
+
+        // Convert Map<Element, List<_>> to Map<Ptr, Integer>
+        for (var el : to) {
+            var ptr = new Ptr(el);
+            var count = refs.get(el).size();
+            referenceCounts.put(ptr, count);
+        }
+
+        // Check if there are any errors in from
+        var fromUri = from.getSourceFile().toUri();
+        for (var err : errors) {
+            if (err.getSource().toUri().equals(fromUri)) hasErrors = true;
+        }
+    }
+
+    boolean needsUpdate(Set<Ptr> signature) {
+        if (hasErrors) return true;
+        for (var expected : referenceCounts.keySet()) {
+            if (!signature.contains(expected)) return true;
+        }
+        return false;
+    }
+
+    int count(Ptr to) {
+        return referenceCounts.getOrDefault(to, 0);
+    }
+}
diff --git a/src/main/java/org/javacs/JavaLanguageServer.java b/src/main/java/org/javacs/JavaLanguageServer.java
index 86fb7cc..8ac317c 100644
--- a/src/main/java/org/javacs/JavaLanguageServer.java
+++ b/src/main/java/org/javacs/JavaLanguageServer.java
@@ -816,9 +816,9 @@ class JavaLanguageServer extends LanguageServer {
         }
 
         // Compile all files that *might* contain references to toEl
-        var fromFiles = compiler.potentialReferences(toEl.get());
-        fromFiles.add(toUri);
-        var batch = compiler.compileBatch(pruneWord(fromFiles, toEl.get()));
+        var fromUris = compiler.potentialReferences(toEl.get());
+        fromUris.add(toUri);
+        var batch = compiler.compileBatch(pruneWord(fromUris, toEl.get()));
 
         // Find toEl again, so that we have an Element from the current batch
         var toElAgain = batch.element(toUri, toLine, toColumn).get();
@@ -985,7 +985,7 @@ class JavaLanguageServer extends LanguageServer {
                 var lens = new CodeLens(range.get(), command, null);
                 result.add(lens);
             }
-            if (!cacheParse.isTestMethod(d)) {
+            if (!cacheParse.isTestMethod(d) && !cacheParse.isTestClass(d)) {
                 // Unresolved "_ references" code lens
                 var start = range.get().start;
                 var line = start.line;
@@ -1045,98 +1045,83 @@ class JavaLanguageServer extends LanguageServer {
             return -1;
         }
 
+        // TODO if new Ptr(toEl) has already been counted, we don't need to re-count unless from-files have been
+        // modified or contain errors
+
         // Compile all files that *might* contain references to toEl
-        var fromFiles = compiler.potentialReferences(toEl.get());
-        fromFiles.add(toUri);
+        var fromUris = compiler.potentialReferences(toEl.get());
+        fromUris.add(toUri);
 
         // If it's too expensive to compute the code lens
-        if (fromFiles.size() > 10) return 100;
+        if (fromUris.size() > 10) return 100;
 
-        // Make sure all fromFiles -> toUri references are in the cache
-        updateCountReferencesCache(toUri, fromFiles);
+        // Make sure all fromUris -> toUri references are in the cache
+        var declarations = activeFileCache.declarations();
+        updateCountReferencesCache(fromUris, toUri, declarations);
 
-        // Count up how many total references exist in fromFiles
+        // Count up how many total references exist in fromUris
         var toPtr = new Ptr(toEl.get());
         var count = 0;
-        for (var from : fromFiles) {
-            var cachedFileCounts = countReferencesCache.get(from);
-            count += cachedFileCounts.counts.getOrDefault(toPtr, 0);
+        for (var fromUri : fromUris) {
+            var fromFile = Paths.get(fromUri);
+            var toFile = Paths.get(toUri);
+            var cachedFileCounts = countReferencesCache.get(fromFile, toFile);
+            count += cachedFileCounts.count(toPtr);
         }
         return 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();
-    }
+    /** countReferencesCache.get(fromFile, toFile) is a count of all references from fromFile to toFile */
+    private final Cache<Path, Index> countReferencesCache = new Cache<>();
 
-    /** countReferencesCacheFile is the file pointed to by every ptr in countReferencesCache[_][ptr] */
+    /**
+     * countReferencesCacheFile is the target of every reference currently in countReferencesCache. Index#needsUpdate(_)
+     * assumes that the user edits one file at a time, and checks whether edits to that file invalidate the cached
+     * index. To guarantee this assumption, we simply invalidate all cached indices when the user changes files.
+     */
     private URI countReferencesCacheFile = URI.create("file:///NONE");
 
-    private void updateCountReferencesCache(URI toFile, Collection<URI> fromFiles) {
-        // If cached file has changed, invalidate the whole cache
-        if (!toFile.equals(countReferencesCacheFile)) {
-            LOG.info(String.format("Cache count-references %s", Parser.fileName(toFile)));
+    private void updateCountReferencesCache(Collection<URI> fromUris, URI toUri, Collection<Element> toEls) {
+        // If the user changes files, invalidate all cached indices
+        if (!toUri.equals(countReferencesCacheFile)) {
             countReferencesCache.clear();
-            countReferencesCacheFile = toFile;
+            countReferencesCacheFile = toUri;
         }
-
-        // 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);
-            }
+        // Convert Element to Ptr, which can be compared across compilation tasks
+        var toPtrs = new HashSet<Ptr>();
+        for (var el : toEls) {
+            toPtrs.add(new Ptr(el));
         }
-
-        // Compile all out-of-date files
-        if (outOfDate.isEmpty()) return;
-        LOG.info(
-                String.format(
-                        "...%d files need to be counted for references to %s",
-                        outOfDate.size(), Parser.fileName(toFile)));
-        outOfDate.add(toFile);
-        countReferencesCache.remove(toFile);
-        var batch = compiler.compileBatch(latestText(outOfDate));
+        // Check which files need to be udpated
+        var outOfDate = new HashSet<URI>();
+        var toFile = Paths.get(toUri);
+        for (var fromUri : fromUris) {
+            var fromFile = Paths.get(fromUri);
+            var needsUpdate =
+                    countReferencesCache.needs(fromFile, toFile)
+                            || countReferencesCache.get(fromFile, toFile).needsUpdate(toPtrs);
+            if (needsUpdate) outOfDate.add(fromFile.toUri());
+        }
+        // If indexes are all up-to-date, we are done
+        if (outOfDate.isEmpty()) {
+            LOG.info(String.format("...all references to %s are already indexed", Parser.fileName(toUri)));
+            return;
+        }
+        // Compile all files that need to be updated in a batch
+        outOfDate.add(toUri);
+        var batch = compiler.compileBatch(outOfDate);
 
         // Find all declarations in toFile
-        var allEls = batch.declarations(toFile);
-
-        // Find all references to all declarations
-        var refs = batch.references(allEls);
+        var toElsAgain = batch.declarations(toUri);
 
-        // Reset all counts for files we just re-compiled
+        // Index outOfDate
+        LOG.info(
+                String.format(
+                        "...search for references to %d elements in %d files", toElsAgain.size(), outOfDate.size()));
         for (var fromUri : outOfDate) {
-            countReferencesCache.put(fromUri, new CountReferences());
-        }
-
-        // Increment 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());
+            var fromFile = Paths.get(fromUri);
+            var index = batch.index(fromUri, toElsAgain);
+            countReferencesCache.load(fromFile, toFile, index);
         }
     }
 
diff --git a/src/main/java/org/javacs/WarnUnused.java b/src/main/java/org/javacs/WarnUnused.java
index dc2646c..614c887 100644
--- a/src/main/java/org/javacs/WarnUnused.java
+++ b/src/main/java/org/javacs/WarnUnused.java
@@ -86,4 +86,10 @@ class WarnUnused extends TreePathScanner<Void, Void> {
         used.add(current());
         return super.visitMemberReference(t, null);
     }
+
+    @Override
+    public Void visitNewClass(NewClassTree t, Void __) {
+        used.add(current());
+        return super.visitNewClass(t, null);
+    }
 }