Re-use compilation tasks

7 files changed, 600 insertions, 184 deletions

diff --git a/src/main/java/org/javacs/CompileBatch.java b/src/main/java/org/javacs/CompileBatch.java
index de6988c..ddfb505 100644
--- a/src/main/java/org/javacs/CompileBatch.java
+++ b/src/main/java/org/javacs/CompileBatch.java
@@ -12,10 +12,10 @@ import javax.lang.model.util.*;
 import javax.tools.*;
 import org.javacs.lsp.Range;
 
-public class CompileBatch {
+public class CompileBatch implements AutoCloseable {
     private final JavaCompilerService parent;
     private final ReportProgress progress;
-    private final JavacTask task;
+    private final TaskPool.Borrow borrow;
     private final Trees trees;
     private final Elements elements;
     private final Types types;
@@ -24,14 +24,14 @@ public class CompileBatch {
     CompileBatch(JavaCompilerService parent, Collection<? extends JavaFileObject> files, ReportProgress progress) {
         this.parent = parent;
         this.progress = progress;
-        this.task = batchTask(parent, files);
-        this.trees = Trees.instance(task);
-        this.elements = task.getElements();
-        this.types = task.getTypes();
+        this.borrow = batchTask(parent, files);
+        this.trees = Trees.instance(borrow.task);
+        this.elements = borrow.task.getElements();
+        this.types = borrow.task.getTypes();
         this.roots = new ArrayList<CompilationUnitTree>();
         // Print timing information for optimization
         var profiler = new Profiler();
-        task.addTaskListener(profiler);
+        borrow.task.addTaskListener(profiler);
         // Show progress message through the UI
         class CountFiles implements TaskListener {
             Set<URI> parse = new HashSet<>(), enter = new HashSet<>(), analyze = new HashSet<>();
@@ -59,35 +59,39 @@ public class CompileBatch {
                 }
             }
         }
-        task.addTaskListener(new CountFiles());
+        borrow.task.addTaskListener(new CountFiles());
         // Compile all roots
         try {
-            for (var t : task.parse()) roots.add(t);
-            // The results of task.analyze() are unreliable when errors are present
+            for (var t : borrow.task.parse()) roots.add(t);
+            // The results of borrow.task.analyze() are unreliable when errors are present
             // You can get at `Element` values using `Trees`
-            task.analyze();
+            borrow.task.analyze();
         } catch (IOException e) {
             throw new RuntimeException(e);
         }
         profiler.print();
     }
 
-    static JavacTask batchTask(JavaCompilerService parent, Collection<? extends JavaFileObject> sources) {
+    @Override
+    public void close() {
+        borrow.close();
+    }
+
+    static TaskPool.Borrow batchTask(JavaCompilerService parent, Collection<? extends JavaFileObject> sources) {
         parent.diags.clear();
-        return (JavacTask)
-                parent.compiler.getTask(
-                        null,
-                        parent.fileManager,
-                        parent.diags::add,
-                        JavaCompilerService.options(parent.classPath),
-                        Collections.emptyList(),
-                        sources);
+        return parent.compiler.getTask(
+                null,
+                parent.fileManager,
+                parent.diags::add,
+                JavaCompilerService.options(parent.classPath),
+                Collections.emptyList(),
+                sources);
     }
 
     public Optional<Element> element(URI uri, int line, int character) {
         for (var root : roots) {
             if (root.getSourceFile().toUri().equals(uri)) {
-                var path = CompileFocus.findPath(task, root, line, character);
+                var path = CompileFocus.findPath(borrow.task, root, line, character);
                 var el = trees.getElement(path);
                 return Optional.ofNullable(el);
             }
@@ -171,7 +175,7 @@ public class CompileBatch {
         // Otherwise, scan roots for references
         List<TreePath> list = new ArrayList<TreePath>();
         var map = Map.of(to, list);
-        var finder = new FindReferences(task);
+        var finder = new FindReferences(borrow.task);
         for (var r : roots) {
             finder.scan(r, map);
         }
@@ -207,7 +211,7 @@ public class CompileBatch {
     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);
+                return new Index(borrow.task, r, parent.diags, declarations);
             }
         }
         throw new RuntimeException(from + " is not in compiled batch");
@@ -216,7 +220,7 @@ public class CompileBatch {
     public Optional<Range> range(TreePath path) {
         var uri = path.getCompilationUnit().getSourceFile().toUri();
         var contents = FileStore.contents(uri);
-        return ParseFile.range(task, contents, path);
+        return ParseFile.range(borrow.task, contents, path);
     }
 
     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
index 46029ae..bbe7b47 100644
--- a/src/main/java/org/javacs/CompileFile.java
+++ b/src/main/java/org/javacs/CompileFile.java
@@ -10,11 +10,11 @@ import java.util.logging.Logger;
 import javax.lang.model.element.*;
 import org.javacs.lsp.*;
 
-public class CompileFile {
+public class CompileFile implements AutoCloseable {
     private final JavaCompilerService parent;
     public final URI file;
     public final String contents;
-    private final JavacTask task;
+    private final TaskPool.Borrow borrow;
     private final Trees trees;
     public final CompilationUnitTree root;
 
@@ -22,21 +22,26 @@ public class CompileFile {
         this.parent = parent;
         this.file = file;
         this.contents = FileStore.contents(file);
-        this.task = CompileFocus.singleFileTask(parent, file, contents);
-        this.trees = Trees.instance(task);
+        this.borrow = CompileFocus.singleFileTask(parent, file, contents);
+        this.trees = Trees.instance(borrow.task);
         var profiler = new Profiler();
-        task.addTaskListener(profiler);
+        borrow.task.addTaskListener(profiler);
         try {
-            this.root = task.parse().iterator().next();
-            // The results of task.analyze() are unreliable when errors are present
+            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`
-            task.analyze();
+            borrow.task.analyze();
         } catch (IOException e) {
             throw new RuntimeException(e);
         }
         profiler.print();
     }
 
+    @Override
+    public void close() {
+        borrow.close();
+    }
+
     public SourcePositions sourcePositions() {
         return trees.getSourcePositions();
     }
@@ -57,14 +62,14 @@ public class CompileFile {
     }
 
     public Index index(List<Element> declarations) {
-        return new Index(task, root, parent.diags, 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(task, root, line, character);
+        var path = CompileFocus.findPath(borrow.task, root, line, character);
         if (path == null) {
             // LOG.info("...found nothing");
             return Optional.empty();
@@ -132,12 +137,12 @@ public class CompileFile {
     }
 
     public Optional<Range> range(TreePath path) {
-        return ParseFile.range(task, contents, path);
+        return ParseFile.range(borrow.task, contents, path);
     }
 
     private List<Element> overrides(ExecutableElement method) {
-        var elements = task.getElements();
-        var types = task.getTypes();
+        var elements = borrow.task.getElements();
+        var types = borrow.task.getTypes();
         var results = new ArrayList<Element>();
         var enclosingClass = (TypeElement) method.getEnclosingElement();
         var enclosingType = enclosingClass.asType();
@@ -219,7 +224,7 @@ public class CompileFile {
         classes.addAll(parent.classPathClasses);
         var fixes = Parser.resolveSymbols(unresolved, sourcePathImports, classes);
         // Figure out which existing imports are actually used
-        var trees = Trees.instance(task);
+        var trees = Trees.instance(borrow.task);
         var references = new HashSet<String>();
         class FindUsedImports extends TreePathScanner<Void, Void> {
             @Override
diff --git a/src/main/java/org/javacs/CompileFocus.java b/src/main/java/org/javacs/CompileFocus.java
index 38a71d4..c1df941 100644
--- a/src/main/java/org/javacs/CompileFocus.java
+++ b/src/main/java/org/javacs/CompileFocus.java
@@ -18,14 +18,14 @@ import javax.lang.model.util.Types;
 import javax.tools.JavaFileObject;
 import javax.tools.StandardLocation;
 
-public class CompileFocus {
+public class CompileFocus implements AutoCloseable {
     public static final int MAX_COMPLETION_ITEMS = 50;
 
     private final JavaCompilerService parent;
     private final URI file;
     private final String contents;
     private final int line, character;
-    private final JavacTask task;
+    private final TaskPool.Borrow borrow;
     private final Trees trees;
     private final Types types;
     private final CompilationUnitTree root;
@@ -37,35 +37,40 @@ public class CompileFocus {
         this.contents = Pruner.prune(file, line, character);
         this.line = line;
         this.character = character;
-        this.task = singleFileTask(parent, file, this.contents);
-        this.trees = Trees.instance(task);
-        this.types = task.getTypes();
+        this.borrow = singleFileTask(parent, file, this.contents);
+        this.trees = Trees.instance(borrow.task);
+        this.types = borrow.task.getTypes();
 
         var profiler = new Profiler();
-        task.addTaskListener(profiler);
+        borrow.task.addTaskListener(profiler);
         try {
-            this.root = task.parse().iterator().next();
+            this.root = borrow.task.parse().iterator().next();
             // The results of task.analyze() are unreliable when errors are present
             // You can get at `Element` values using `Trees`
-            task.analyze();
+            borrow.task.analyze();
         } catch (IOException e) {
             throw new RuntimeException(e);
         }
         profiler.print();
-        this.path = findPath(task, root, line, character);
+        this.path = findPath(borrow.task, root, line, character);
     }
 
+    @Override
+    public void close() {
+        borrow.close();
+    }
+
+    // TODO this is highlighting incorrectly
     /** Create a task that compiles a single file */
-    static JavacTask singleFileTask(JavaCompilerService parent, URI file, String contents) {
+    static TaskPool.Borrow singleFileTask(JavaCompilerService parent, URI file, String contents) {
         parent.diags.clear();
-        return (JavacTask)
-                parent.compiler.getTask(
-                        null,
-                        parent.fileManager,
-                        parent.diags::add,
-                        JavaCompilerService.options(parent.classPath),
-                        Collections.emptyList(),
-                        List.of(new SourceFileObject(file, contents)));
+        return parent.compiler.getTask(
+                null,
+                parent.fileManager,
+                parent.diags::add,
+                JavaCompilerService.options(parent.classPath),
+                Collections.emptyList(),
+                List.of(new SourceFileObject(file, contents)));
     }
 
     /** Find the smallest element that includes the cursor */
@@ -203,8 +208,12 @@ public class CompileFocus {
 
         // Get members of switched type
         var type = trees.getTypeMirror(path);
+        if (type == null) {
+            LOG.info("...no type at " + leaf.getExpression());
+            return Collections.emptyList();
+        }
         LOG.info(String.format("...switched expression has type `%s`", type));
-        var types = task.getTypes();
+        var types = borrow.task.getTypes();
         var definition = types.asElement(type);
         if (definition == null) {
             LOG.info("...type has no definition, completing identifiers instead");
@@ -221,7 +230,7 @@ public class CompileFocus {
 
     /** Find all members of expression ending at line:character */
     public List<Completion> completeMembers(boolean isReference) {
-        var types = task.getTypes();
+        var types = borrow.task.getTypes();
         var scope = trees.getScope(path);
         var element = trees.getElement(path);
 
@@ -426,7 +435,7 @@ public class CompileFocus {
     }
 
     private void collectSuperMethods(TypeMirror thisType, List<ExecutableElement> result) {
-        var types = task.getTypes();
+        var types = borrow.task.getTypes();
 
         for (var superType : types.directSupertypes(thisType)) {
             if (superType instanceof DeclaredType) {
@@ -472,7 +481,7 @@ public class CompileFocus {
         collectSupers(t, types);
         // Object type is not included by default
         // We need to add it to get members like .equals(other) and .hashCode()
-        types.add(task.getElements().getTypeElement("java.lang.Object").asType());
+        types.add(borrow.task.getElements().getTypeElement("java.lang.Object").asType());
         return types;
     }
 
@@ -514,7 +523,7 @@ public class CompileFocus {
 
     /** Find all identifiers in scope at line:character */
     List<Element> scopeMembers(String partialName) {
-        var types = task.getTypes();
+        var types = borrow.task.getTypes();
         var start = trees.getScope(path);
 
         class Walk {
diff --git a/src/main/java/org/javacs/JavaCompilerService.java b/src/main/java/org/javacs/JavaCompilerService.java
index 8dfeaf8..afa0906 100644
--- a/src/main/java/org/javacs/JavaCompilerService.java
+++ b/src/main/java/org/javacs/JavaCompilerService.java
@@ -16,7 +16,7 @@ import javax.tools.*;
 public class JavaCompilerService {
     // Not modifiable! If you want to edit these, you need to create a new instance
     final Set<Path> classPath, docPath;
-    final JavaCompiler compiler = ServiceLoader.load(JavaCompiler.class).iterator().next();
+    final TaskPool compiler = new TaskPool(10);
     final Docs docs;
     final Set<String> jdkClasses = Classes.jdkTopLevelClasses(), classPathClasses;
     // Diagnostics from the last compilation task
@@ -40,7 +40,6 @@ public class JavaCompilerService {
         this.docs = new Docs(docPath);
         this.classPathClasses = Classes.classPathTopLevelClasses(classPath);
         this.fileManager = new SourceFileManager();
-        ;
     }
 
     /** Combine source path or class path entries using the system separator, for example ':' in unix */
@@ -121,46 +120,46 @@ public class JavaCompilerService {
 
         // Create task
         var options = options(classPath);
-        var task =
-                (JavacTask) compiler.getTask(null, fileManager, diags::add, options, Collections.emptyList(), sources);
-        var trees = Trees.instance(task);
-
-        // Print timing information for optimization
-        var profiler = new Profiler();
-        task.addTaskListener(profiler);
-
-        // Run compilation
-        diags.clear();
-        Iterable<? extends CompilationUnitTree> roots;
-        try {
-            roots = task.parse();
-            task.analyze();
-        } catch (IOException e) {
-            throw new RuntimeException(e);
-        }
-        profiler.print();
-        LOG.info(String.format("...found %d errors", diags.size()));
-
-        // Check for unused privates
-        for (var r : roots) {
-            var warnUnused = new WarnUnused(task);
-            warnUnused.scan(r, null);
-            for (var unusedEl : warnUnused.notUsed()) {
-                var path = trees.getPath(unusedEl);
-                var message = String.format("`%s` is not used", unusedEl.getSimpleName());
-                Diagnostic.Kind kind;
-                if (unusedEl instanceof ExecutableElement || unusedEl instanceof TypeElement) {
-                    kind = Diagnostic.Kind.OTHER;
-                } else {
-                    kind = Diagnostic.Kind.WARNING;
+        try (var borrow = compiler.getTask(null, fileManager, diags::add, options, Collections.emptyList(), sources)) {
+            var trees = Trees.instance(borrow.task);
+
+            // Print timing information for optimization
+            var profiler = new Profiler();
+            borrow.task.addTaskListener(profiler);
+
+            // Run compilation
+            diags.clear();
+            Iterable<? extends CompilationUnitTree> roots;
+            try {
+                roots = borrow.task.parse();
+                borrow.task.analyze();
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+            profiler.print();
+            LOG.info(String.format("...found %d errors", diags.size()));
+
+            // Check for unused privates
+            for (var r : roots) {
+                var warnUnused = new WarnUnused(borrow.task);
+                warnUnused.scan(r, null);
+                for (var unusedEl : warnUnused.notUsed()) {
+                    var path = trees.getPath(unusedEl);
+                    var message = String.format("`%s` is not used", unusedEl.getSimpleName());
+                    Diagnostic.Kind kind;
+                    if (unusedEl instanceof ExecutableElement || unusedEl instanceof TypeElement) {
+                        kind = Diagnostic.Kind.OTHER;
+                    } else {
+                        kind = Diagnostic.Kind.WARNING;
+                    }
+                    diags.add(new Warning(borrow.task, path, kind, "unused", message));
                 }
-                diags.add(new Warning(task, path, kind, "unused", message));
             }
-        }
-        // TODO hint fields that could be final
-        // TODO hint unused exception
+            // TODO hint fields that could be final
+            // TODO hint unused exception
 
-        return Collections.unmodifiableList(new ArrayList<>(diags));
+            return Collections.unmodifiableList(new ArrayList<>(diags));
+        }
     }
 
     public Set<URI> potentialDefinitions(Element to) {
diff --git a/src/main/java/org/javacs/JavaLanguageServer.java b/src/main/java/org/javacs/JavaLanguageServer.java
index 780d8e5..b1d9894 100644
--- a/src/main/java/org/javacs/JavaLanguageServer.java
+++ b/src/main/java/org/javacs/JavaLanguageServer.java
@@ -351,35 +351,34 @@ class JavaLanguageServer extends LanguageServer {
         }
         // Compile again, focusing on a region that depends on what type of completion we want to do
         var ctx = maybeCtx.get();
-        // TODO CompileFocus should have a "patch" mechanism where we recompile the current file without creating a new
-        // task
-        var focus = compiler.compileFocus(uri, ctx.line, ctx.character);
-        // Do a specific type of completion
         List<Completion> cs;
         boolean isIncomplete;
-        switch (ctx.kind) {
-            case MemberSelect:
-                cs = focus.completeMembers(false);
-                isIncomplete = false;
-                break;
-            case MemberReference:
-                cs = focus.completeMembers(true);
-                isIncomplete = false;
-                break;
-            case Identifier:
-                cs = focus.completeIdentifiers(ctx.inClass, ctx.inMethod, ctx.partialName);
-                isIncomplete = cs.size() >= CompileFocus.MAX_COMPLETION_ITEMS;
-                break;
-            case Annotation:
-                cs = focus.completeAnnotations(ctx.partialName);
-                isIncomplete = cs.size() >= CompileFocus.MAX_COMPLETION_ITEMS;
-                break;
-            case Case:
-                cs = focus.completeCases();
-                isIncomplete = false;
-                break;
-            default:
-                throw new RuntimeException("Unexpected completion context " + ctx.kind);
+        try (var focus = compiler.compileFocus(uri, ctx.line, ctx.character)) {
+            // Do a specific type of completion
+            switch (ctx.kind) {
+                case MemberSelect:
+                    cs = focus.completeMembers(false);
+                    isIncomplete = false;
+                    break;
+                case MemberReference:
+                    cs = focus.completeMembers(true);
+                    isIncomplete = false;
+                    break;
+                case Identifier:
+                    cs = focus.completeIdentifiers(ctx.inClass, ctx.inMethod, ctx.partialName);
+                    isIncomplete = cs.size() >= CompileFocus.MAX_COMPLETION_ITEMS;
+                    break;
+                case Annotation:
+                    cs = focus.completeAnnotations(ctx.partialName);
+                    isIncomplete = cs.size() >= CompileFocus.MAX_COMPLETION_ITEMS;
+                    break;
+                case Case:
+                    cs = focus.completeCases();
+                    isIncomplete = false;
+                    break;
+                default:
+                    throw new RuntimeException("Unexpected completion context " + ctx.kind);
+            }
         }
         // Convert to CompletionItem
         var result = new ArrayList<CompletionItem>();
@@ -624,6 +623,7 @@ class JavaLanguageServer extends LanguageServer {
                 || !activeFileCache.file.equals(uri)
                 || activeFileCacheVersion != FileStore.version(uri)) {
             LOG.info("Recompile active file...");
+            if (activeFileCache != null) activeFileCache.close();
             activeFileCache = compiler.compileFile(uri);
             activeFileCacheVersion = FileStore.version(uri);
         }
@@ -744,11 +744,10 @@ class JavaLanguageServer extends LanguageServer {
         if (!FileStore.isJavaFile(uri)) return Optional.empty();
         var line = position.position.line + 1;
         var column = position.position.character + 1;
-        // TODO CompileFocus should have a "patch" mechanism where we recompile the current file without creating a new
-        // task
-        var focus = compiler.compileFocus(uri, line, column);
-        var help = focus.methodInvocation().map(this::asSignatureHelp);
-        return help;
+        try (var focus = compiler.compileFocus(uri, line, column)) {
+            var help = focus.methodInvocation().map(this::asSignatureHelp);
+            return help;
+        }
     }
 
     @Override
@@ -770,26 +769,26 @@ class JavaLanguageServer extends LanguageServer {
         // Compile all files that *might* contain definitions of fromEl
         var toFiles = compiler.potentialDefinitions(toEl.get());
         toFiles.add(fromUri);
-        var batch = compiler.compileBatch(pruneWord(toFiles, toEl.get()));
-
-        // 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.isPresent()) return Optional.empty();
-        var result = new ArrayList<Location>();
-        for (var path : toTreePaths.get()) {
-            var toUri = path.getCompilationUnit().getSourceFile().toUri();
-            var toRange = batch.range(path);
-            if (!toRange.isPresent()) {
-                LOG.warning(String.format("Couldn't locate `%s`", path.getLeaf()));
-                continue;
+        try (var batch = compiler.compileBatch(pruneWord(toFiles, toEl.get()))) {
+            // 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.isPresent()) return Optional.empty();
+            var result = new ArrayList<Location>();
+            for (var path : toTreePaths.get()) {
+                var toUri = path.getCompilationUnit().getSourceFile().toUri();
+                var toRange = batch.range(path);
+                if (!toRange.isPresent()) {
+                    LOG.warning(String.format("Couldn't locate `%s`", path.getLeaf()));
+                    continue;
+                }
+                var from = new Location(toUri, toRange.get());
+                result.add(from);
             }
-            var from = new Location(toUri, toRange.get());
-            result.add(from);
+            return Optional.of(result);
         }
-        return Optional.of(result);
     }
 
     @Override
@@ -811,26 +810,26 @@ class JavaLanguageServer extends LanguageServer {
         // Compile all files that *might* contain references to toEl
         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();
-
-        // Find all references to toElAgain
-        var fromTreePaths = batch.references(toElAgain);
-        if (!fromTreePaths.isPresent()) return Optional.empty();
-        var result = new ArrayList<Location>();
-        for (var path : fromTreePaths.get()) {
-            var fromUri = path.getCompilationUnit().getSourceFile().toUri();
-            var fromRange = batch.range(path);
-            if (!fromRange.isPresent()) {
-                LOG.warning(String.format("Couldn't locate `%s`", path.getLeaf()));
-                continue;
+        try (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();
+
+            // Find all references to toElAgain
+            var fromTreePaths = batch.references(toElAgain);
+            if (!fromTreePaths.isPresent()) return Optional.empty();
+            var result = new ArrayList<Location>();
+            for (var path : fromTreePaths.get()) {
+                var fromUri = path.getCompilationUnit().getSourceFile().toUri();
+                var fromRange = batch.range(path);
+                if (!fromRange.isPresent()) {
+                    LOG.warning(String.format("Couldn't locate `%s`", path.getLeaf()));
+                    continue;
+                }
+                var from = new Location(fromUri, fromRange.get());
+                result.add(from);
             }
-            var from = new Location(fromUri, fromRange.get());
-            result.add(from);
+            return Optional.of(result);
         }
-        return Optional.of(result);
     }
 
     private List<JavaFileObject> pruneWord(Collection<URI> files, Element el) {
@@ -1138,18 +1137,18 @@ class JavaLanguageServer extends LanguageServer {
         if (!outOfDate.isEmpty()) {
             // 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 toEls = batch.declarations(toUri);
+            try (var batch = compiler.compileBatch(outOfDate)) {
+                // Find all declarations in toFile
+                var toEls = batch.declarations(toUri);
 
-            // Index outOfDate
-            LOG.info(
-                    String.format(
-                            "...search for references to %d elements in %d files", toEls.size(), outOfDate.size()));
-            for (var fromUri : outOfDate) {
-                var index = batch.index(fromUri, toEls);
-                cacheIndex.put(fromUri, index);
+                // Index outOfDate
+                LOG.info(
+                        String.format(
+                                "...search for references to %d elements in %d files", toEls.size(), outOfDate.size()));
+                for (var fromUri : outOfDate) {
+                    var index = batch.index(fromUri, toEls);
+                    cacheIndex.put(fromUri, index);
+                }
             }
         } else {
             LOG.info("...all indexes are cached and up-to-date");
diff --git a/src/main/java/org/javacs/ParseFile.java b/src/main/java/org/javacs/ParseFile.java
index 33521ed..fe83637 100644
--- a/src/main/java/org/javacs/ParseFile.java
+++ b/src/main/java/org/javacs/ParseFile.java
@@ -5,15 +5,29 @@ import com.sun.source.tree.*;
 import com.sun.source.util.*;
 import java.io.IOException;
 import java.net.URI;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Objects;
-import java.util.Optional;
+import java.util.*;
 import java.util.logging.Logger;
 import javax.lang.model.element.*;
+import javax.tools.JavaCompiler;
 import org.javacs.lsp.*;
 
 public class ParseFile {
+    private static final JavaCompiler COMPILER = ServiceLoader.load(JavaCompiler.class).iterator().next();
+
+    /** Create a task that compiles a single file */
+    static JavacTask singleFileTask(JavaCompilerService parent, URI file, String contents) {
+        // TODO could eliminate the connection to parent
+        parent.diags.clear();
+        return (JavacTask)
+                COMPILER.getTask(
+                        null,
+                        parent.fileManager,
+                        parent.diags::add,
+                        JavaCompilerService.options(parent.classPath),
+                        Collections.emptyList(),
+                        List.of(new SourceFileObject(file, contents)));
+    }
+
     private final String contents;
     private final JavacTask task;
     private final Trees trees;
@@ -24,7 +38,7 @@ public class ParseFile {
         Objects.requireNonNull(file);
 
         this.contents = FileStore.contents(file);
-        this.task = CompileFocus.singleFileTask(parent, file, contents);
+        this.task = singleFileTask(parent, file, contents);
         this.trees = Trees.instance(task);
         var profiler = new Profiler();
         task.addTaskListener(profiler);
diff --git a/src/main/java/org/javacs/TaskPool.java b/src/main/java/org/javacs/TaskPool.java
new file mode 100644
index 0000000..8475ae9
--- /dev/null
+++ b/src/main/java/org/javacs/TaskPool.java
@@ -0,0 +1,386 @@
+// Forked from JavacTaskImpl
+/*
+ * Copyright (c) 2015, 2017, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.  Oracle designates this
+ * particular file as subject to the "Classpath" exception as provided
+ * by Oracle in the LICENSE file that accompanied this code.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+package org.javacs;
+
+import com.sun.source.tree.CompilationUnitTree;
+import com.sun.source.util.JavacTask;
+import com.sun.source.util.TaskEvent;
+import com.sun.source.util.TaskEvent.Kind;
+import com.sun.source.util.TaskListener;
+import com.sun.tools.javac.api.*;
+import com.sun.tools.javac.code.Symtab;
+import com.sun.tools.javac.code.Types;
+import com.sun.tools.javac.comp.Annotate;
+import com.sun.tools.javac.comp.Check;
+import com.sun.tools.javac.comp.CompileStates;
+import com.sun.tools.javac.comp.Enter;
+import com.sun.tools.javac.comp.Modules;
+import com.sun.tools.javac.main.Arguments;
+import com.sun.tools.javac.main.JavaCompiler;
+import com.sun.tools.javac.model.JavacElements;
+import com.sun.tools.javac.util.Context;
+import com.sun.tools.javac.util.DefinedBy;
+import com.sun.tools.javac.util.DefinedBy.Api;
+import com.sun.tools.javac.util.Log;
+import java.io.PrintStream;
+import java.io.Writer;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+import javax.tools.Diagnostic;
+import javax.tools.DiagnosticListener;
+import javax.tools.JavaFileManager;
+import javax.tools.JavaFileObject;
+
+/**
+ * A pool of reusable JavacTasks. When a task is no valid anymore, it is returned to the pool, and its Context may be
+ * reused for future processing in some cases. The reuse is achieved by replacing some components (most notably
+ * JavaCompiler and Log) with reusable counterparts, and by cleaning up leftovers from previous compilation.
+ *
+ * <p>For each combination of options, a separate task/context is created and kept, as most option values are cached
+ * inside components themselves.
+ *
+ * <p>When the compilation redefines sensitive classes (e.g. classes in the the java.* packages), the task/context is
+ * not reused.
+ *
+ * <p>When the task is reused, then packages that were already listed won't be listed again.
+ *
+ * <p>Care must be taken to only return tasks that won't be used by the original caller.
+ *
+ * <p>Care must also be taken when custom components are installed, as those are not cleaned when the task/context is
+ * reused, and subsequent getTask may return a task based on a context with these custom components.
+ *
+ * <p><b>This is NOT part of any supported API. If you write code that depends on this, you do so at your own risk. This
+ * code and its internal interfaces are subject to change or deletion without notice.</b>
+ */
+public class TaskPool {
+
+    private static final JavacTool systemProvider = JavacTool.create();
+
+    private final int maxPoolSize;
+    private final Map<List<String>, List<ReusableContext>> options2Contexts = new HashMap<>();
+    private int id;
+
+    private int statReused = 0;
+    private int statNew = 0;
+    private int statPolluted = 0;
+    private int statRemoved = 0;
+
+    /**
+     * Creates the pool.
+     *
+     * @param maxPoolSize maximum number of tasks/context that will be kept in the pool.
+     */
+    public TaskPool(int maxPoolSize) {
+        this.maxPoolSize = maxPoolSize;
+    }
+
+    /**
+     * Creates a new task as if by {@link javax.tools.JavaCompiler#getTask} and runs the provided worker with it. The
+     * task is only valid while the worker is running. The internal structures may be reused from some previous
+     * compilation.
+     *
+     * @param out a Writer for additional output from the compiler; use {@code System.err} if {@code null}
+     * @param fileManager a file manager; if {@code null} use the compiler's standard filemanager
+     * @param diagnosticListener a diagnostic listener; if {@code null} use the compiler's default method for reporting
+     *     diagnostics
+     * @param options compiler options, {@code null} means no options
+     * @param classes names of classes to be processed by annotation processing, {@code null} means no class names
+     * @param compilationUnits the compilation units to compile, {@code null} means no compilation units
+     * @param worker that should be run with the task
+     * @return an object representing the compilation
+     * @throws RuntimeException if an unrecoverable error occurred in a user supplied component. The {@linkplain
+     *     Throwable#getCause() cause} will be the error in user code.
+     * @throws IllegalArgumentException if any of the options are invalid, or if any of the given compilation units are
+     *     of other kind than {@linkplain JavaFileObject.Kind#SOURCE source}
+     */
+    public Borrow getTask(
+            Writer out,
+            JavaFileManager fileManager,
+            DiagnosticListener<? super JavaFileObject> diagnosticListener,
+            Iterable<String> options,
+            Iterable<String> classes,
+            Iterable<? extends JavaFileObject> compilationUnits) {
+        List<String> opts =
+                StreamSupport.stream(options.spliterator(), false).collect(Collectors.toCollection(ArrayList::new));
+
+        ReusableContext ctx;
+
+        synchronized (this) {
+            List<ReusableContext> cached = options2Contexts.getOrDefault(opts, Collections.emptyList());
+
+            if (cached.isEmpty()) {
+                ctx = new ReusableContext(opts);
+                statNew++;
+            } else {
+                ctx = cached.remove(0);
+                statReused++;
+            }
+        }
+
+        ctx.useCount++;
+
+        JavacTaskImpl task =
+                (JavacTaskImpl)
+                        systemProvider.getTask(
+                                out, fileManager, diagnosticListener, opts, classes, compilationUnits, ctx);
+
+        task.addTaskListener(ctx);
+
+        return new Borrow(task, ctx);
+    }
+
+    // where:
+    private long cacheSize() {
+        return options2Contexts.values().stream().flatMap(Collection::stream).count();
+    }
+
+    public void printStatistics(PrintStream out) {
+        out.println(statReused + " reused Contexts");
+        out.println(statNew + " newly created Contexts");
+        out.println(statPolluted + " polluted Contexts");
+        out.println(statRemoved + " removed Contexts");
+    }
+
+    public class Borrow implements AutoCloseable {
+        public final JavacTask task;
+        public final ReusableContext ctx;
+
+        public Borrow(JavacTask task, ReusableContext ctx) {
+            this.task = task;
+            this.ctx = ctx;
+        }
+
+        @Override
+        public void close() {
+            // not returning the context to the pool if task crashes with an exception
+            // the task/context may be in a broken state
+            ctx.clear();
+            if (ctx.polluted) {
+                statPolluted++;
+            } else {
+                //            task.cleanup();
+                synchronized (this) {
+                    while (cacheSize() + 1 > maxPoolSize) {
+                        ReusableContext toRemove =
+                                options2Contexts
+                                        .values()
+                                        .stream()
+                                        .flatMap(Collection::stream)
+                                        .sorted((c1, c2) -> c1.timeStamp < c2.timeStamp ? -1 : 1)
+                                        .findFirst()
+                                        .get();
+                        options2Contexts.get(toRemove.arguments).remove(toRemove);
+                        statRemoved++;
+                    }
+                    options2Contexts.computeIfAbsent(ctx.arguments, x -> new ArrayList<>()).add(ctx);
+                    ctx.timeStamp = id++;
+                }
+            }
+        }
+    }
+
+    static class ReusableContext extends Context implements TaskListener {
+
+        Set<CompilationUnitTree> roots = new HashSet<>();
+
+        List<String> arguments;
+        boolean polluted = false;
+
+        int useCount;
+        long timeStamp;
+
+        ReusableContext(List<String> arguments) {
+            super();
+            this.arguments = arguments;
+            put(Log.logKey, ReusableLog.factory);
+            put(JavaCompiler.compilerKey, ReusableJavaCompiler.factory);
+        }
+
+        void clear() {
+            drop(Arguments.argsKey);
+            drop(DiagnosticListener.class);
+            drop(Log.outKey);
+            drop(Log.errKey);
+            drop(JavaFileManager.class);
+            drop(JavacTask.class);
+            drop(JavacTrees.class);
+            drop(JavacElements.class);
+
+            if (ht.get(Log.logKey) instanceof ReusableLog) {
+                // log already inited - not first round
+                ((ReusableLog) Log.instance(this)).clear();
+                Enter.instance(this).newRound();
+                ((ReusableJavaCompiler) ReusableJavaCompiler.instance(this)).clear();
+                Types.instance(this).newRound();
+                Check.instance(this).newRound();
+                Modules.instance(this).newRound();
+                Annotate.instance(this).newRound();
+                CompileStates.instance(this).clear();
+                MultiTaskListener.instance(this).clear();
+
+                // find if any of the roots have redefined java.* classes
+                Symtab syms = Symtab.instance(this);
+                //                pollutionScanner.scan(roots, syms);
+                roots.clear();
+            }
+        }
+
+        /**
+         * This scanner detects as to whether the shared context has been polluted. This happens whenever a compiled
+         * program redefines a core class (in 'java.*' package) or when (typically because of cyclic inheritance) the
+         * symbol kind of a core class has been touched.
+         */
+        /*
+        TreeScanner<Void, Symtab> pollutionScanner = new TreeScanner<Void, Symtab>() {
+            @Override @DefinedBy(Api.COMPILER_TREE)
+            public Void visitClass(ClassTree node, Symtab syms) {
+                Symbol sym = ((JCClassDecl)node).sym;
+                if (sym != null) {
+                    syms.removeClass(sym.packge().modle, sym.flatName());
+                    Type sup = supertype(sym);
+                    if (isCoreClass(sym) ||
+                            (sup != null && isCoreClass(sup.tsym) && sup.tsym.kind != Kinds.Kind.TYP)) {
+                        polluted = true;
+                    }
+                }
+                return super.visitClass(node, syms);
+            }
+
+            private boolean isCoreClass(Symbol s) {
+                return s.flatName().toString().startsWith("java.");
+            }
+
+            private Type supertype(Symbol s) {
+                if (s.type == null ||
+                        !s.type.hasTag(TypeTag.CLASS)) {
+                    return null;
+                } else {
+                    ClassType ct = (ClassType)s.type;
+                    return ct.supertype_field;
+                }
+            }
+        };
+        */
+
+        @Override
+        @DefinedBy(Api.COMPILER_TREE)
+        public void finished(TaskEvent e) {
+            if (e.getKind() == Kind.PARSE) {
+                roots.add(e.getCompilationUnit());
+            }
+        }
+
+        @Override
+        @DefinedBy(Api.COMPILER_TREE)
+        public void started(TaskEvent e) {
+            // do nothing
+        }
+
+        <T> void drop(Key<T> k) {
+            ht.remove(k);
+        }
+
+        <T> void drop(Class<T> c) {
+            ht.remove(key(c));
+        }
+
+        /**
+         * Reusable JavaCompiler; exposes a method to clean up the component from leftovers associated with previous
+         * compilations.
+         */
+        static class ReusableJavaCompiler extends JavaCompiler {
+
+            static final Factory<JavaCompiler> factory = ReusableJavaCompiler::new;
+
+            ReusableJavaCompiler(Context context) {
+                super(context);
+            }
+
+            @Override
+            public void close() {
+                // do nothing
+            }
+
+            void clear() {
+                newRound();
+            }
+
+            @Override
+            protected void checkReusable() {
+                // do nothing - it's ok to reuse the compiler
+            }
+        }
+
+        /**
+         * Reusable Log; exposes a method to clean up the component from leftovers associated with previous
+         * compilations.
+         */
+        static class ReusableLog extends Log {
+
+            static final Factory<Log> factory = ReusableLog::new;
+
+            Context context;
+
+            ReusableLog(Context context) {
+                super(context);
+                this.context = context;
+            }
+
+            void clear() {
+                recorded.clear();
+                sourceMap.clear();
+                nerrors = 0;
+                nwarnings = 0;
+                // Set a fake listener that will lazily lookup the context for the 'real' listener. Since
+                // this field is never updated when a new task is created, we cannot simply reset the field
+                // or keep old value. This is a hack to workaround the limitations in the current infrastructure.
+                diagListener =
+                        new DiagnosticListener<JavaFileObject>() {
+                            DiagnosticListener<JavaFileObject> cachedListener;
+
+                            @Override
+                            @DefinedBy(Api.COMPILER)
+                            @SuppressWarnings("unchecked")
+                            public void report(Diagnostic<? extends JavaFileObject> diagnostic) {
+                                if (cachedListener == null) {
+                                    cachedListener = context.get(DiagnosticListener.class);
+                                }
+                                cachedListener.report(diagnostic);
+                            }
+                        };
+            }
+        }
+    }
+}