Always update active file

4 files changed, 80 insertions, 40 deletions

diff --git a/src/main/java/org/javacs/CompileBatch.java b/src/main/java/org/javacs/CompileBatch.java
index 47d66fa..947ba84 100644
--- a/src/main/java/org/javacs/CompileBatch.java
+++ b/src/main/java/org/javacs/CompileBatch.java
@@ -179,15 +179,6 @@ public class CompileBatch {
         return Optional.of(list);
     }
 
-    public Index index(URI from, List<Element> toAny) {
-        for (var r : roots) {
-            if (r.getSourceFile().toUri().equals(from)) {
-                return new Index(task, r, parent.diags, toAny);
-            }
-        }
-        throw new RuntimeException(from + " is not in compiled batch");
-    }
-
     /**
      * Find all elements in `file` that get turned into code-lenses. This needs to match the result of
      * `ParseFile#declarations`
@@ -213,6 +204,15 @@ public class CompileBatch {
         throw new RuntimeException(file + " is not in " + message);
     }
 
+    public Index index(URI from, List<Element> declarations) {
+        for (var r : roots) {
+            if (r.getSourceFile().toUri().equals(from)) {
+                return new Index(task, r, parent.diags, declarations);
+            }
+        }
+        throw new RuntimeException(from + " is not in compiled batch");
+    }
+
     public Optional<Range> range(TreePath path) {
         var uri = path.getCompilationUnit().getSourceFile().toUri();
         var contents = FileStore.contents(uri);
diff --git a/src/main/java/org/javacs/CompileFile.java b/src/main/java/org/javacs/CompileFile.java
index 8bcef9c..7269edc 100644
--- a/src/main/java/org/javacs/CompileFile.java
+++ b/src/main/java/org/javacs/CompileFile.java
@@ -62,6 +62,10 @@ public class CompileFile {
         return els;
     }
 
+    public Index index(List<Element> declarations) {
+        return new Index(task, root, parent.diags, declarations);
+    }
+
     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/Index.java b/src/main/java/org/javacs/Index.java
index 6ab404d..66bb822 100644
--- a/src/main/java/org/javacs/Index.java
+++ b/src/main/java/org/javacs/Index.java
@@ -15,7 +15,7 @@ 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;
+    final boolean hasErrors;
 
     Index(
             JavacTask task,
@@ -38,15 +38,25 @@ class Index {
         }
 
         // Check if there are any errors in from
+        this.hasErrors = hasErrors(from, errors);
+    }
+
+    private static boolean hasErrors(
+            CompilationUnitTree from, Collection<Diagnostic<? extends JavaFileObject>> errors) {
         var fromUri = from.getSourceFile().toUri();
         for (var err : errors) {
-            if (err.getSource().toUri().equals(fromUri)) hasErrors = true;
+            if (err.getSource().toUri().equals(fromUri) && err.getKind() == Diagnostic.Kind.ERROR) {
+                return true;
+            }
         }
+        return false;
     }
 
     boolean needsUpdate(Set<Ptr> signature) {
         if (hasErrors) return true;
         for (var expected : referenceCounts.keySet()) {
+            // Note: in theory, you could change the return type of a method, and this could secretly change the
+            // references counts
             if (!signature.contains(expected)) return true;
         }
         return false;
@@ -55,4 +65,12 @@ class Index {
     int count(Ptr to) {
         return referenceCounts.getOrDefault(to, 0);
     }
+
+    int total() {
+        var sum = 0;
+        for (var count : referenceCounts.values()) {
+            sum += count;
+        }
+        return sum;
+    }
 }
diff --git a/src/main/java/org/javacs/JavaLanguageServer.java b/src/main/java/org/javacs/JavaLanguageServer.java
index f641922..c280064 100644
--- a/src/main/java/org/javacs/JavaLanguageServer.java
+++ b/src/main/java/org/javacs/JavaLanguageServer.java
@@ -1036,25 +1036,25 @@ class JavaLanguageServer extends LanguageServer {
     }
 
     /**
-     * countReferencesCcacheCountReferencesFileacheFile 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.
+     * cacheReferencesFile is the target of every reference currently in cacheIndex. 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 cacheCountReferencesFile = URI.create("file:///NONE");
+    private URI cacheReferencesFile = URI.create("file:///NONE");
     /** cacheCountReferences[toDeclaration] is a list of all files that have references to toDeclaration */
-    private final Map<Ptr, List<Index>> cacheCountReferences = new HashMap<>();
+    private final Map<Ptr, List<URI>> cacheReferences = new HashMap<>();
     /**
      * cacheCountReferences[toDeclaration] == TOO_EXPENSIVE indicates there are too many potential references to
      * toDeclaration
      */
-    private static final List<Index> TOO_EXPENSIVE = new ArrayList<>();
-    /** countReferencesCache[fromFile] is a count of all references from fromFile to cacheCountReferencesFile */
+    private static final List<URI> TOO_EXPENSIVE = new ArrayList<>();
+    /** cacheIndex[fromFile] is a count of all references from fromFile to cacheCountReferencesFile */
     private final Map<URI, Index> cacheIndex = new HashMap<>();
 
-    private boolean cacheCountReferencesNeedsUpdate(Ptr toPtr, Set<Ptr> signature) {
-        if (!cacheCountReferences.containsKey(toPtr)) return true;
-        for (var index : cacheCountReferences.get(toPtr)) {
+    private boolean cacheReferencesNeedsUpdate(Ptr toPtr, Set<Ptr> signature) {
+        if (!cacheReferences.containsKey(toPtr)) return true;
+        for (var fromUri : cacheReferences.get(toPtr)) {
+            var index = cacheIndex.get(fromUri);
             if (index.needsUpdate(signature)) return true;
         }
         return false;
@@ -1062,10 +1062,10 @@ class JavaLanguageServer extends LanguageServer {
 
     private int countReferences(URI toUri, int toLine, int toColumn) {
         // If the user changes files, invalidate all cached indices
-        if (!toUri.equals(cacheCountReferencesFile)) {
-            cacheCountReferences.clear();
+        if (!toUri.equals(cacheReferencesFile)) {
+            cacheReferences.clear();
             cacheIndex.clear();
-            cacheCountReferencesFile = toUri;
+            cacheReferencesFile = toUri;
         }
 
         // Make sure the active file is compiled
@@ -1087,12 +1087,12 @@ class JavaLanguageServer extends LanguageServer {
         }
 
         // If the signature has changed, or the from-files contain errors, we need to redo the count
-        if (cacheCountReferencesNeedsUpdate(toPtr, signature)) {
+        if (cacheReferencesNeedsUpdate(toPtr, signature)) {
             LOG.info(String.format("Count references to `%s`...", toPtr));
 
             // Compile all files that *might* contain references to toEl
             var fromUris = compiler.potentialReferences(toEl.get());
-            fromUris.add(toUri);
+            fromUris.remove(toUri);
 
             // If it's too expensive to compute the code lens
             if (fromUris.size() > 10) {
@@ -1100,21 +1100,25 @@ class JavaLanguageServer extends LanguageServer {
                         String.format(
                                 "...there are %d potential references, which is too expensive to compile",
                                 fromUris.size()));
-                cacheCountReferences.put(toPtr, TOO_EXPENSIVE);
+                cacheReferences.put(toPtr, TOO_EXPENSIVE);
             } else {
                 // Make sure all fromUris -> toUri references are in cacheIndex
-                var list = index(fromUris, toUri, signature);
-                cacheCountReferences.put(toPtr, list);
+                var list = referencesFile(fromUris, toUri, signature);
+                cacheReferences.put(toPtr, list);
             }
         } else {
             LOG.info(String.format("Using cached count references to `%s`", toPtr));
         }
 
+        // Always update active file
+        var activeIndex = activeFileCache.index(declarations);
+        var count = activeIndex.count(toPtr);
+
         // Count up references out of index
-        var count = 0;
-        var list = cacheCountReferences.get(toPtr);
-        if (list == TOO_EXPENSIVE) return 100;
-        for (var index : list) {
+        var fromUris = cacheReferences.get(toPtr);
+        if (fromUris == TOO_EXPENSIVE) return 100;
+        for (var fromUri : fromUris) {
+            var index = cacheIndex.get(fromUri);
             count += index.count(toPtr);
         }
 
@@ -1123,10 +1127,23 @@ class JavaLanguageServer extends LanguageServer {
 
     private boolean cacheIndexNeedsUpdate(URI fromUri, Set<Ptr> signature) {
         if (!cacheIndex.containsKey(fromUri)) return true;
-        return cacheIndex.get(fromUri).needsUpdate(signature);
+        var index = cacheIndex.get(fromUri);
+        if (index.hasErrors) {
+            LOG.info(
+                    String.format("...%s needs to be re-indexed because it contains errors", Parser.fileName(fromUri)));
+            return true;
+        }
+        if (index.needsUpdate(signature)) {
+            LOG.info(
+                    String.format(
+                            "...%s needs to be re-indexed because it refers to a declaration that has changed",
+                            Parser.fileName(fromUri)));
+            return true;
+        }
+        return false;
     }
 
-    private List<Index> index(Collection<URI> fromUris, URI toUri, Set<Ptr> signature) {
+    private List<URI> referencesFile(Collection<URI> fromUris, URI toUri, Set<Ptr> signature) {
         // Check which files need to be updated
         var outOfDate = new HashSet<URI>();
         for (var fromUri : fromUris) {
@@ -1156,12 +1173,13 @@ class JavaLanguageServer extends LanguageServer {
             LOG.info("...all indexes are cached and up-to-date");
         }
 
-        // List all indices
-        var list = new ArrayList<Index>();
+        // Figure out which files actually reference one of the Ptrs in signature
+        var actuallyReferencesFile = new ArrayList<URI>();
         for (var fromUri : fromUris) {
-            list.add(cacheIndex.get(fromUri));
+            var index = cacheIndex.get(fromUri);
+            if (index.total() > 0) actuallyReferencesFile.add(fromUri);
         }
-        return list;
+        return actuallyReferencesFile;
     }
 
     @Override