Combine CompileFile and CompileBatch

4 files changed, 186 insertions, 245 deletions

diff --git a/src/main/java/org/javacs/CompileBatch.java b/src/main/java/org/javacs/CompileBatch.java
index ddfb505..f12c486 100644
--- a/src/main/java/org/javacs/CompileBatch.java
+++ b/src/main/java/org/javacs/CompileBatch.java
@@ -88,12 +88,10 @@ public class CompileBatch implements AutoCloseable {
                 sources);
     }
 
-    public Optional<Element> element(URI uri, int line, int character) {
+    public CompilationUnitTree root(URI uri) {
         for (var root : roots) {
             if (root.getSourceFile().toUri().equals(uri)) {
-                var path = CompileFocus.findPath(borrow.task, root, line, character);
-                var el = trees.getElement(path);
-                return Optional.ofNullable(el);
+                return root;
             }
         }
         // Somehow, uri was not in batch
@@ -104,6 +102,21 @@ public class CompileBatch implements AutoCloseable {
         throw new RuntimeException("File " + uri + " isn't in batch " + names);
     }
 
+    private String contents(CompilationUnitTree root) {
+        try {
+            return root.getSourceFile().getCharContent(true).toString();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public Optional<Element> element(URI uri, int line, int character) {
+        var root = root(uri);
+        var path = CompileFocus.findPath(borrow.task, root, line, character);
+        var el = trees.getElement(path);
+        return Optional.ofNullable(el);
+    }
+
     public Optional<List<TreePath>> definitions(Element el) {
         LOG.info(String.format("Search for definitions of `%s` in %d files...", el, roots.size()));
 
@@ -223,5 +236,150 @@ public class CompileBatch implements AutoCloseable {
         return ParseFile.range(borrow.task, contents, path);
     }
 
+    public SourcePositions sourcePositions() {
+        return trees.getSourcePositions();
+    }
+
+    public LineMap lineMap(URI uri) {
+        return root(uri).getLineMap();
+    }
+
+    public List<? extends ImportTree> imports(URI uri) {
+        return root(uri).getImports();
+    }
+
+    private List<Element> overrides(ExecutableElement method) {
+        var elements = borrow.task.getElements();
+        var types = borrow.task.getTypes();
+        var results = new ArrayList<Element>();
+        var enclosingClass = (TypeElement) method.getEnclosingElement();
+        var enclosingType = enclosingClass.asType();
+        for (var superClass : types.directSupertypes(enclosingType)) {
+            var e = (TypeElement) types.asElement(superClass);
+            for (var other : e.getEnclosedElements()) {
+                if (!(other instanceof ExecutableElement)) continue;
+                if (elements.overrides(method, (ExecutableElement) other, enclosingClass)) {
+                    results.add(other);
+                }
+            }
+        }
+        return results;
+    }
+
+    private boolean hasOverrideAnnotation(ExecutableElement method) {
+        for (var ann : method.getAnnotationMirrors()) {
+            var type = ann.getAnnotationType();
+            var el = type.asElement();
+            var name = el.toString();
+            if (name.equals("java.lang.Override")) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    /** Find methods that override a method from a superclass but don't have an @Override annotation. */
+    public List<TreePath> needsOverrideAnnotation(URI uri) {
+        LOG.info(String.format("Looking for methods that need an @Override annotation in %s ...", uri.getPath()));
+
+        var root = root(uri);
+        var results = new ArrayList<TreePath>();
+        class FindMissingOverride extends TreePathScanner<Void, Void> {
+            @Override
+            public Void visitMethod(MethodTree t, Void __) {
+                var method = (ExecutableElement) trees.getElement(getCurrentPath());
+                var supers = overrides(method);
+                if (!supers.isEmpty() && !hasOverrideAnnotation(method)) {
+                    var overridesMethod = supers.get(0);
+                    var overridesClass = overridesMethod.getEnclosingElement();
+                    LOG.info(
+                            String.format(
+                                    "...`%s` has no @Override annotation but overrides `%s.%s`",
+                                    method, overridesClass, overridesMethod));
+                    results.add(getCurrentPath());
+                }
+                return super.visitMethod(t, null);
+            }
+        }
+        new FindMissingOverride().scan(root, null);
+        return results;
+    }
+
+    /**
+     * Figure out what imports this file should have. Star-imports like `import java.util.*` are converted to individual
+     * class imports. Missing imports are inferred by looking at imports in other source files.
+     */
+    public List<String> fixImports(URI uri) {
+        var root = root(uri);
+        var contents = contents(root);
+        // Check diagnostics for missing imports
+        var unresolved = new HashSet<String>();
+        for (var d : parent.diags) {
+            if (d.getCode().equals("compiler.err.cant.resolve.location") && d.getSource().toUri().equals(uri)) {
+                long start = d.getStartPosition(), end = d.getEndPosition();
+                var id = contents.substring((int) start, (int) end);
+                if (id.matches("[A-Z]\\w+")) {
+                    unresolved.add(id);
+                } else LOG.warning(id + " doesn't look like a class");
+            } else if (d.getMessage(null).contains("cannot find to")) {
+                var lines = d.getMessage(null).split("\n");
+                var firstLine = lines.length > 0 ? lines[0] : "";
+                LOG.warning(String.format("%s %s doesn't look like to-not-found", d.getCode(), firstLine));
+            }
+        }
+        // Look at imports in other classes to help us guess how to fix imports
+        // TODO cache parsed imports on a per-file basis
+        var sourcePathImports = Parser.existingImports(FileStore.all());
+        var classes = new HashSet<String>();
+        classes.addAll(parent.jdkClasses);
+        classes.addAll(parent.classPathClasses);
+        var fixes = Parser.resolveSymbols(unresolved, sourcePathImports, classes);
+        // Figure out which existing imports are actually used
+        var trees = Trees.instance(borrow.task);
+        var references = new HashSet<String>();
+        class FindUsedImports extends TreePathScanner<Void, Void> {
+            @Override
+            public Void visitIdentifier(IdentifierTree node, Void nothing) {
+                var e = trees.getElement(getCurrentPath());
+                if (e instanceof TypeElement) {
+                    var t = (TypeElement) e;
+                    var qualifiedName = t.getQualifiedName().toString();
+                    var lastDot = qualifiedName.lastIndexOf('.');
+                    var packageName = lastDot == -1 ? "" : qualifiedName.substring(0, lastDot);
+                    var thisPackage = Objects.toString(root.getPackageName(), "");
+                    // java.lang.* and current package are imported by default
+                    if (!packageName.equals("java.lang")
+                            && !packageName.equals(thisPackage)
+                            && !packageName.equals("")) {
+                        references.add(qualifiedName);
+                    }
+                }
+                return null;
+            }
+        }
+        new FindUsedImports().scan(root, null);
+        // Take the intersection of existing imports ^ existing identifiers
+        var qualifiedNames = new HashSet<String>();
+        for (var i : root.getImports()) {
+            var imported = i.getQualifiedIdentifier().toString();
+            if (imported.endsWith(".*")) {
+                var packageName = Parser.mostName(imported);
+                var isUsed = references.stream().anyMatch(r -> r.startsWith(packageName));
+                if (isUsed) qualifiedNames.add(imported);
+                else LOG.warning("There are no references to package " + imported);
+            } else {
+                if (references.contains(imported)) qualifiedNames.add(imported);
+                else LOG.warning("There are no references to class " + imported);
+            }
+        }
+        // Add qualified names from fixes
+        qualifiedNames.addAll(fixes.values());
+        // Sort in alphabetical order
+        var sorted = new ArrayList<String>();
+        sorted.addAll(qualifiedNames);
+        Collections.sort(sorted);
+        return sorted;
+    }
+
     private static final Logger LOG = Logger.getLogger("main");
 }
diff --git a/src/main/java/org/javacs/CompileFile.java b/src/main/java/org/javacs/CompileFile.java
deleted file mode 100644
index 7d00270..0000000
--- a/src/main/java/org/javacs/CompileFile.java
+++ /dev/null
@@ -1,218 +0,0 @@
-package org.javacs;
-
-import com.sun.source.tree.*;
-import com.sun.source.util.*;
-import java.io.IOException;
-import java.net.URI;
-import java.util.*;
-import java.util.logging.Logger;
-import javax.lang.model.element.*;
-
-public class CompileFile implements AutoCloseable {
-    private final JavaCompilerService parent;
-    public final URI file;
-    public final String contents;
-    private final TaskPool.Borrow borrow;
-    private final Trees trees;
-    public final CompilationUnitTree root;
-
-    CompileFile(JavaCompilerService parent, URI file) {
-        this.parent = parent;
-        this.file = file;
-        this.contents = FileStore.contents(file);
-        this.borrow = CompileFocus.singleFileTask(parent, file, contents);
-        this.trees = Trees.instance(borrow.task);
-        var profiler = new Profiler();
-        borrow.task.addTaskListener(profiler);
-        try {
-            this.root = borrow.task.parse().iterator().next();
-            // The results of borrow.task.analyze() are unreliable when errors are present
-            // You can get at `Element` values using `Trees`
-            borrow.task.analyze();
-        } catch (IOException e) {
-            throw new RuntimeException(e);
-        }
-        profiler.print();
-    }
-
-    @Override
-    public void close() {
-        borrow.close();
-    }
-
-    public SourcePositions sourcePositions() {
-        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 Index index(List<Element> declarations) {
-        return new Index(borrow.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));
-
-        // First, look for a tree path
-        var path = CompileFocus.findPath(borrow.task, root, line, character);
-        if (path == null) {
-            // LOG.info("...found nothing");
-            return Optional.empty();
-        }
-        // LOG.info(String.format("...found tree `%s`", Parser.describeTree(path.getLeaf())));
-
-        // Then, convert the path to an element
-        var el = trees.getElement(path);
-        if (el == null) {
-            // LOG.info(String.format("...tree does not correspond to an element"));
-            return Optional.empty();
-        }
-
-        return Optional.of(el);
-    }
-
-    private List<Element> overrides(ExecutableElement method) {
-        var elements = borrow.task.getElements();
-        var types = borrow.task.getTypes();
-        var results = new ArrayList<Element>();
-        var enclosingClass = (TypeElement) method.getEnclosingElement();
-        var enclosingType = enclosingClass.asType();
-        for (var superClass : types.directSupertypes(enclosingType)) {
-            var e = (TypeElement) types.asElement(superClass);
-            for (var other : e.getEnclosedElements()) {
-                if (!(other instanceof ExecutableElement)) continue;
-                if (elements.overrides(method, (ExecutableElement) other, enclosingClass)) {
-                    results.add(other);
-                }
-            }
-        }
-        return results;
-    }
-
-    private boolean hasOverrideAnnotation(ExecutableElement method) {
-        for (var ann : method.getAnnotationMirrors()) {
-            var type = ann.getAnnotationType();
-            var el = type.asElement();
-            var name = el.toString();
-            if (name.equals("java.lang.Override")) {
-                return true;
-            }
-        }
-        return false;
-    }
-
-    /** Find methods that override a method from a superclass but don't have an @Override annotation. */
-    public List<TreePath> needsOverrideAnnotation() {
-        LOG.info(String.format("Looking for methods that need an @Override annotation in %s ...", file.getPath()));
-
-        var results = new ArrayList<TreePath>();
-        class FindMissingOverride extends TreePathScanner<Void, Void> {
-            @Override
-            public Void visitMethod(MethodTree t, Void __) {
-                var method = (ExecutableElement) trees.getElement(getCurrentPath());
-                var supers = overrides(method);
-                if (!supers.isEmpty() && !hasOverrideAnnotation(method)) {
-                    var overridesMethod = supers.get(0);
-                    var overridesClass = overridesMethod.getEnclosingElement();
-                    LOG.info(
-                            String.format(
-                                    "...`%s` has no @Override annotation but overrides `%s.%s`",
-                                    method, overridesClass, overridesMethod));
-                    results.add(getCurrentPath());
-                }
-                return super.visitMethod(t, null);
-            }
-        }
-        new FindMissingOverride().scan(root, null);
-        return results;
-    }
-
-    /**
-     * Figure out what imports this file should have. Star-imports like `import java.util.*` are converted to individual
-     * class imports. Missing imports are inferred by looking at imports in other source files.
-     */
-    public List<String> fixImports() {
-        // Check diagnostics for missing imports
-        var unresolved = new HashSet<String>();
-        for (var d : parent.diags) {
-            if (d.getCode().equals("compiler.err.cant.resolve.location") && d.getSource().toUri().equals(file)) {
-                long start = d.getStartPosition(), end = d.getEndPosition();
-                var id = contents.substring((int) start, (int) end);
-                if (id.matches("[A-Z]\\w+")) {
-                    unresolved.add(id);
-                } else LOG.warning(id + " doesn't look like a class");
-            } else if (d.getMessage(null).contains("cannot find to")) {
-                var lines = d.getMessage(null).split("\n");
-                var firstLine = lines.length > 0 ? lines[0] : "";
-                LOG.warning(String.format("%s %s doesn't look like to-not-found", d.getCode(), firstLine));
-            }
-        }
-        // Look at imports in other classes to help us guess how to fix imports
-        // TODO cache parsed imports on a per-file basis
-        var sourcePathImports = Parser.existingImports(FileStore.all());
-        var classes = new HashSet<String>();
-        classes.addAll(parent.jdkClasses);
-        classes.addAll(parent.classPathClasses);
-        var fixes = Parser.resolveSymbols(unresolved, sourcePathImports, classes);
-        // Figure out which existing imports are actually used
-        var trees = Trees.instance(borrow.task);
-        var references = new HashSet<String>();
-        class FindUsedImports extends TreePathScanner<Void, Void> {
-            @Override
-            public Void visitIdentifier(IdentifierTree node, Void nothing) {
-                var e = trees.getElement(getCurrentPath());
-                if (e instanceof TypeElement) {
-                    var t = (TypeElement) e;
-                    var qualifiedName = t.getQualifiedName().toString();
-                    var lastDot = qualifiedName.lastIndexOf('.');
-                    var packageName = lastDot == -1 ? "" : qualifiedName.substring(0, lastDot);
-                    var thisPackage = Objects.toString(root.getPackageName(), "");
-                    // java.lang.* and current package are imported by default
-                    if (!packageName.equals("java.lang")
-                            && !packageName.equals(thisPackage)
-                            && !packageName.equals("")) {
-                        references.add(qualifiedName);
-                    }
-                }
-                return null;
-            }
-        }
-        new FindUsedImports().scan(root, null);
-        // Take the intersection of existing imports ^ existing identifiers
-        var qualifiedNames = new HashSet<String>();
-        for (var i : root.getImports()) {
-            var imported = i.getQualifiedIdentifier().toString();
-            if (imported.endsWith(".*")) {
-                var packageName = Parser.mostName(imported);
-                var isUsed = references.stream().anyMatch(r -> r.startsWith(packageName));
-                if (isUsed) qualifiedNames.add(imported);
-                else LOG.warning("There are no references to package " + imported);
-            } else {
-                if (references.contains(imported)) qualifiedNames.add(imported);
-                else LOG.warning("There are no references to class " + imported);
-            }
-        }
-        // Add qualified names from fixes
-        qualifiedNames.addAll(fixes.values());
-        // Sort in alphabetical order
-        var sorted = new ArrayList<String>();
-        sorted.addAll(qualifiedNames);
-        Collections.sort(sorted);
-        return sorted;
-    }
-
-    private static final Logger LOG = Logger.getLogger("main");
-}
diff --git a/src/main/java/org/javacs/JavaCompilerService.java b/src/main/java/org/javacs/JavaCompilerService.java
index afa0906..871fefb 100644
--- a/src/main/java/org/javacs/JavaCompilerService.java
+++ b/src/main/java/org/javacs/JavaCompilerService.java
@@ -83,10 +83,6 @@ public class JavaCompilerService {
         return new CompileFocus(this, file, line, character);
     }
 
-    public CompileFile compileFile(URI file) {
-        return new CompileFile(this, file);
-    }
-
     public CompileBatch compileBatch(Collection<URI> uris) {
         return compileBatch(uris, ReportProgress.EMPTY);
     }
diff --git a/src/main/java/org/javacs/JavaLanguageServer.java b/src/main/java/org/javacs/JavaLanguageServer.java
index b1d9894..c66b61b 100644
--- a/src/main/java/org/javacs/JavaLanguageServer.java
+++ b/src/main/java/org/javacs/JavaLanguageServer.java
@@ -614,17 +614,19 @@ class JavaLanguageServer extends LanguageServer {
         return Optional.of(md);
     }
 
-    private CompileFile activeFileCache;
+    private CompileBatch activeFileCache;
+    private URI activeFileCacheFile = URI.create("file:///NONE");
     private int activeFileCacheVersion = -1;
 
     // TODO take only URI and invalidate based on version
     private void updateActiveFile(URI uri) {
         if (activeFileCache == null
-                || !activeFileCache.file.equals(uri)
+                || !activeFileCacheFile.equals(uri)
                 || activeFileCacheVersion != FileStore.version(uri)) {
             LOG.info("Recompile active file...");
             if (activeFileCache != null) activeFileCache.close();
-            activeFileCache = compiler.compileFile(uri);
+            activeFileCache = compiler.compileBatch(Collections.singleton(uri));
+            activeFileCacheFile = uri;
             activeFileCacheVersion = FileStore.version(uri);
         }
     }
@@ -639,7 +641,7 @@ class JavaLanguageServer extends LanguageServer {
         // Find element undeer cursor
         var line = position.position.line + 1;
         var column = position.position.character + 1;
-        var el = activeFileCache.element(line, column);
+        var el = activeFileCache.element(uri, line, column);
         if (!el.isPresent()) return Optional.empty();
 
         var result = new ArrayList<MarkedString>();
@@ -760,7 +762,7 @@ class JavaLanguageServer extends LanguageServer {
         // Compile from-file and identify element under cursor
         LOG.info(String.format("Go-to-def at %s:%d...", fromUri, fromLine));
         updateActiveFile(fromUri);
-        var toEl = activeFileCache.element(fromLine, fromColumn);
+        var toEl = activeFileCache.element(fromUri, fromLine, fromColumn);
         if (!toEl.isPresent()) {
             LOG.info(String.format("...no element at cursor"));
             return Optional.empty();
@@ -801,7 +803,7 @@ class JavaLanguageServer extends LanguageServer {
         // Compile from-file and identify element under cursor
         LOG.warning(String.format("Looking for references to %s(%d,%d)...", toUri.getPath(), toLine, toColumn));
         updateActiveFile(toUri);
-        var toEl = activeFileCache.element(toLine, toColumn);
+        var toEl = activeFileCache.element(toUri, toLine, toColumn);
         if (!toEl.isPresent()) {
             LOG.warning("...no element under cursor");
             return Optional.empty();
@@ -1053,7 +1055,7 @@ class JavaLanguageServer extends LanguageServer {
         updateActiveFile(toUri);
 
         // Find the element we want to count references to
-        var toEl = activeFileCache.element(toLine, toColumn);
+        var toEl = activeFileCache.element(toUri, toLine, toColumn);
         if (!toEl.isPresent()) {
             LOG.warning("...no element at code lens");
             return -1;
@@ -1061,7 +1063,7 @@ class JavaLanguageServer extends LanguageServer {
         var toPtr = new Ptr(toEl.get());
 
         // Find the signature of the target file
-        var declarations = activeFileCache.declarations();
+        var declarations = activeFileCache.declarations(toUri);
         var signature = new HashSet<Ptr>();
         for (var el : declarations) {
             signature.add(new Ptr(el));
@@ -1092,7 +1094,7 @@ class JavaLanguageServer extends LanguageServer {
         }
 
         // Always update active file
-        var activeIndex = activeFileCache.index(declarations);
+        var activeIndex = activeFileCache.index(toUri, declarations);
         var count = activeIndex.count(toPtr);
 
         // Count up references out of index
@@ -1178,14 +1180,14 @@ class JavaLanguageServer extends LanguageServer {
     private List<TextEdit> fixImports() {
         // TODO if imports already match fixed-imports, return empty list
         // TODO preserve comments and other details of existing imports
-        var imports = activeFileCache.fixImports();
+        var imports = activeFileCache.fixImports(activeFileCacheFile);
         var pos = activeFileCache.sourcePositions();
-        var lines = activeFileCache.root.getLineMap();
+        var lines = activeFileCache.lineMap(activeFileCacheFile);
         var edits = new ArrayList<TextEdit>();
         // Delete all existing imports
-        for (var i : activeFileCache.root.getImports()) {
+        for (var i : activeFileCache.imports(activeFileCacheFile)) {
             if (!i.isStatic()) {
-                var offset = pos.getStartPosition(activeFileCache.root, i);
+                var offset = pos.getStartPosition(activeFileCache.root(activeFileCacheFile), i);
                 var line = (int) lines.getLineNumber(offset) - 1;
                 var delete = new TextEdit(new Range(new Position(line, 0), new Position(line + 1, 0)), "");
                 edits.add(delete);
@@ -1196,15 +1198,18 @@ class JavaLanguageServer extends LanguageServer {
         long insertLine = -1;
         var insertText = new StringBuilder();
         // If there are imports, use the start of the first import as the insert position
-        for (var i : activeFileCache.root.getImports()) {
+        for (var i : activeFileCache.imports(activeFileCacheFile)) {
             if (!i.isStatic() && insertLine == -1) {
-                long offset = pos.getStartPosition(activeFileCache.root, i);
+                long offset = pos.getStartPosition(activeFileCache.root(activeFileCacheFile), i);
                 insertLine = lines.getLineNumber(offset) - 1;
             }
         }
         // If there are no imports, insert after the package declaration
-        if (insertLine == -1 && activeFileCache.root.getPackageName() != null) {
-            long offset = pos.getEndPosition(activeFileCache.root, activeFileCache.root.getPackageName());
+        if (insertLine == -1 && activeFileCache.root(activeFileCacheFile).getPackageName() != null) {
+            long offset =
+                    pos.getEndPosition(
+                            activeFileCache.root(activeFileCacheFile),
+                            activeFileCache.root(activeFileCacheFile).getPackageName());
             insertLine = lines.getLineNumber(offset);
             insertText.append("\n");
         }
@@ -1225,9 +1230,9 @@ class JavaLanguageServer extends LanguageServer {
 
     private List<TextEdit> addOverrides() {
         var edits = new ArrayList<TextEdit>();
-        var methods = activeFileCache.needsOverrideAnnotation();
+        var methods = activeFileCache.needsOverrideAnnotation(activeFileCacheFile);
         var pos = activeFileCache.sourcePositions();
-        var lines = activeFileCache.root.getLineMap();
+        var lines = activeFileCache.lineMap(activeFileCacheFile);
         for (var t : methods) {
             var methodStart = pos.getStartPosition(t.getCompilationUnit(), t.getLeaf());
             var insertLine = lines.getLineNumber(methodStart);