Override JavaFileManager more aggressively

4 files changed, 545 insertions, 4 deletions

diff --git a/src/main/java/org/javacs/JarFileObject.java b/src/main/java/org/javacs/JarFileObject.java
new file mode 100644
index 0000000..a88d01c
--- /dev/null
+++ b/src/main/java/org/javacs/JarFileObject.java
@@ -0,0 +1,101 @@
+package org.javacs;
+
+import java.io.*;
+import java.net.URI;
+import java.nio.charset.Charset;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
+import javax.lang.model.element.Modifier;
+import javax.lang.model.element.NestingKind;
+import javax.tools.JavaFileObject;
+
+class JarFileObject implements JavaFileObject {
+    final Path jar;
+    final JarEntry entry;
+
+    JarFileObject(Path jar, JarEntry entry) {
+        this.jar = jar;
+        this.entry = entry;
+    }
+
+    @Override
+    public Kind getKind() {
+        var name = entry.getName().toString();
+        return SourceFileObject.kindFromExtension(name);
+    }
+
+    @Override
+    public boolean isNameCompatible(String simpleName, Kind kind) {
+        var relative = Paths.get(simpleName.replace('.', '/'));
+        var absolute = Paths.get(entry.getName());
+        return absolute.endsWith(relative);
+    }
+
+    @Override
+    public NestingKind getNestingKind() {
+        return null;
+    }
+
+    @Override
+    public Modifier getAccessLevel() {
+        return null;
+    }
+
+    @Override
+    public URI toUri() {
+        var entryName = entry.getName();
+        var jarURI = jar.toUri().normalize();
+        var separator = entryName.startsWith("/") ? "!" : "!/";
+        return URI.create("jar:" + jarURI + separator + entryName);
+    }
+
+    @Override
+    public String getName() {
+        return jar + "(" + entry.getName() + ")";
+    }
+
+    @Override
+    public InputStream openInputStream() throws IOException {
+        return new JarFile(jar.toFile()).getInputStream(entry);
+    }
+
+    @Override
+    public OutputStream openOutputStream() throws IOException {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public Reader openReader(boolean ignoreEncodingErrors) throws IOException {
+        return new InputStreamReader(openInputStream());
+    }
+
+    @Override
+    public CharSequence getCharContent(boolean ignoreEncodingErrors) throws IOException {
+        var bytes = openInputStream().readAllBytes();
+        return new String(bytes, Charset.forName("UTF-8"));
+    }
+
+    @Override
+    public Writer openWriter() throws IOException {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public long getLastModified() {
+        return entry.getLastModifiedTime().toMillis();
+    }
+
+    @Override
+    public boolean delete() {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        if (other.getClass() != JarFileObject.class) return false;
+        var that = (JarFileObject) other;
+        return this.jar.equals(that.jar) && this.entry.getName().equals(that.entry.getName());
+    }
+}
diff --git a/src/main/java/org/javacs/JavaCompilerService.java b/src/main/java/org/javacs/JavaCompilerService.java
index 2d380a1..80a1f4d 100644
--- a/src/main/java/org/javacs/JavaCompilerService.java
+++ b/src/main/java/org/javacs/JavaCompilerService.java
@@ -27,8 +27,8 @@ public class JavaCompilerService {
     final List<Diagnostic<? extends JavaFileObject>> diags = new ArrayList<>();
     // Use the same file manager for multiple tasks, so we don't repeatedly re-compile the same files
     // TODO intercept files that aren't in the batch and erase method bodies so compilation is faster
-    final StandardJavaFileManager fileManager =
-            new FileManagerWrapper(compiler.getStandardFileManager(diags::add, null, Charset.defaultCharset()));
+    final StandardJavaFileManager fileManager;
+    static final boolean useSourceFileManager = true;
 
     public JavaCompilerService(
             Set<Path> sourcePath, Supplier<Set<Path>> allJavaFiles, Set<Path> classPath, Set<Path> docPath) {
@@ -54,6 +54,12 @@ public class JavaCompilerService {
         docSourcePath.addAll(docPath);
         this.docs = new Docs(docSourcePath);
         this.classPathClasses = Classes.classPathTopLevelClasses(classPath);
+        this.fileManager =
+                useSourceFileManager
+                        ? new SourceFileManager(sourcePath, classPath)
+                        : new FileManagerWrapper(
+                                compiler.getStandardFileManager(diags::add, null, Charset.defaultCharset()));
+        ;
     }
 
     /** Combine source path or class path entries using the system separator, for example ':' in unix */
@@ -64,8 +70,10 @@ public class JavaCompilerService {
     static List<String> options(Set<Path> sourcePath, Set<Path> classPath) {
         var list = new ArrayList<String>();
 
-        Collections.addAll(list, "-classpath", joinPath(classPath));
-        Collections.addAll(list, "-sourcepath", joinPath(sourcePath));
+        if (!useSourceFileManager) {
+            Collections.addAll(list, "-classpath", joinPath(classPath));
+            Collections.addAll(list, "-sourcepath", joinPath(sourcePath));
+        }
         // Collections.addAll(list, "-verbose");
         Collections.addAll(list, "-proc:none");
         Collections.addAll(list, "-g");
diff --git a/src/main/java/org/javacs/SourceFileManager.java b/src/main/java/org/javacs/SourceFileManager.java
new file mode 100644
index 0000000..0d8b792
--- /dev/null
+++ b/src/main/java/org/javacs/SourceFileManager.java
@@ -0,0 +1,327 @@
+package org.javacs;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.nio.charset.Charset;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.*;
+import java.util.function.Predicate;
+import java.util.jar.JarFile;
+import java.util.stream.Stream;
+import javax.tools.*;
+
+class SourceFileManager implements StandardJavaFileManager {
+    private final StandardJavaFileManager delegate = createDelegateFileManager();
+
+    private final Set<Path> sourcePath, classPath;
+
+    SourceFileManager(Set<Path> sourcePath, Set<Path> classPath) {
+        this.sourcePath = sourcePath;
+        this.classPath = classPath;
+    }
+
+    private static StandardJavaFileManager createDelegateFileManager() {
+        var compiler = ServiceLoader.load(JavaCompiler.class).iterator().next();
+        return compiler.getStandardFileManager(__ -> {}, null, Charset.defaultCharset());
+    }
+
+    @Override
+    public ClassLoader getClassLoader(Location location) {
+        var thisClassLoader = getClass().getClassLoader();
+        if (location == StandardLocation.CLASS_PATH) {
+            return new URLClassLoader(urls(classPath), thisClassLoader);
+        } else if (location == StandardLocation.SOURCE_PATH) {
+            return new URLClassLoader(urls(sourcePath), thisClassLoader);
+        } else {
+            return thisClassLoader;
+        }
+    }
+
+    private URL[] urls(Set<Path> paths) {
+        var urls = new URL[paths.size()];
+        var i = 0;
+        for (var p : paths) {
+            try {
+                urls[i++] = p.toUri().toURL();
+            } catch (MalformedURLException e) {
+                throw new RuntimeException(e);
+            }
+        }
+        return urls;
+    }
+
+    @Override
+    public Iterable<JavaFileObject> list(
+            Location location, String packageName, Set<JavaFileObject.Kind> kinds, boolean recurse) throws IOException {
+        if (location == StandardLocation.CLASS_PATH) {
+            return list(classPath, this::isClassPathFile);
+        } else if (location == StandardLocation.SOURCE_PATH) {
+            return list(sourcePath, this::isJavaSource);
+        } else {
+            return delegate.list(location, packageName, kinds, recurse);
+        }
+    }
+
+    private boolean isJavaSource(JavaFileObject file) {
+        return file.getName().endsWith(".java");
+    }
+
+    private boolean isClassPathFile(JavaFileObject file) {
+        var name = file.getName();
+        return name.endsWith(".class") || name.endsWith(".jar");
+    }
+
+    private Iterable<JavaFileObject> list(Set<Path> dirs, Predicate<JavaFileObject> filter) {
+        var stream = dirs.stream().flatMap(this::list).flatMap(this::asJavaFileObjects).filter(filter);
+        return stream::iterator;
+    }
+
+    private Stream<Path> list(Path dir) {
+        try {
+            if (!Files.exists(dir)) return Stream.of();
+            return Files.walk(dir).filter(Files::isRegularFile);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private Stream<JavaFileObject> asJavaFileObjects(Path file) {
+        var fileName = file.getFileName().toString();
+        if (fileName.endsWith(".jar")) {
+            JarFile jar;
+            try {
+                jar = new JarFile(file.toFile());
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+            return jar.stream().map(entry -> new JarFileObject(file, entry));
+        } else {
+            return Stream.of(new SourceFileObject(file));
+        }
+    }
+
+    @Override
+    public String inferBinaryName(Location location, JavaFileObject file) {
+        if (location == StandardLocation.SOURCE_PATH) {
+            var source = (SourceFileObject) file;
+            return sourceFileBinaryName(sourcePath, source);
+        } else if (location == StandardLocation.CLASS_PATH && file instanceof SourceFileObject) {
+            var source = (SourceFileObject) file;
+            return sourceFileBinaryName(classPath, source);
+        } else if (location == StandardLocation.CLASS_PATH && file instanceof JarFileObject) {
+            var jarFile = (JarFileObject) file;
+            var relativePath = jarFile.entry.getName();
+            return binaryName(relativePath);
+        } else {
+            return delegate.inferBinaryName(location, file);
+        }
+    }
+
+    private String sourceFileBinaryName(Set<Path> path, SourceFileObject source) {
+        for (var root : path) {
+            if (source.path.startsWith(root)) {
+                var relativePath = root.relativize(source.path).toString();
+                return binaryName(relativePath);
+            }
+        }
+        return null;
+    }
+
+    private String binaryName(String relativePath) {
+        var slash = removeExtension(relativePath);
+        return slash.replace('/', '.');
+    }
+
+    private String removeExtension(String fileName) {
+        var lastDot = fileName.lastIndexOf(".");
+        return (lastDot == -1 ? fileName : fileName.substring(0, lastDot));
+    }
+
+    @Override
+    public boolean isSameFile(FileObject a, FileObject b) {
+        return a.equals(b);
+    }
+
+    @Override
+    public boolean handleOption(String current, Iterator<String> remaining) {
+        return false;
+    }
+
+    @Override
+    public boolean hasLocation(Location location) {
+        return location == StandardLocation.CLASS_PATH
+                || location == StandardLocation.SOURCE_PATH
+                || delegate.hasLocation(location);
+    }
+
+    @Override
+    public JavaFileObject getJavaFileForInput(Location location, String className, JavaFileObject.Kind kind)
+            throws IOException {
+        if (location == StandardLocation.SOURCE_PATH || location == StandardLocation.CLASS_PATH) {
+            var relative = className.replace('.', '/') + kind.extension;
+            return findFileForInput(location, relative);
+        }
+        return delegate.getJavaFileForInput(location, className, kind);
+    }
+
+    @Override
+    public FileObject getFileForInput(Location location, String packageName, String relativeName) throws IOException {
+        if (location == StandardLocation.SOURCE_PATH || location == StandardLocation.CLASS_PATH) {
+            var relative = relativeName;
+            if (!packageName.isEmpty()) {
+                relative = packageName.replace('.', '/') + '/' + relative;
+            }
+            return findFileForInput(location, relative);
+        }
+        return delegate.getFileForInput(location, packageName, relativeName);
+    }
+
+    private JavaFileObject findFileForInput(Location location, String relative) {
+        if (location == StandardLocation.CLASS_PATH) {
+            for (var root : classPath) {
+                if (Files.isDirectory(root)) {
+                    var absolute = root.resolve(relative);
+                    if (Files.exists(absolute)) {
+                        return new SourceFileObject(root);
+                    }
+                } else if (root.getFileName().toString().endsWith(".jar")) {
+                    JarFile jar;
+                    try {
+                        jar = new JarFile(root.toFile());
+                    } catch (IOException e) {
+                        throw new RuntimeException(e);
+                    }
+                    var entry = jar.getJarEntry(relative);
+                    if (entry != null) {
+                        return new JarFileObject(root, entry);
+                    }
+                }
+            }
+        } else if (location == StandardLocation.SOURCE_PATH) {
+            for (var root : sourcePath) {
+                var absolute = root.resolve(relative);
+                if (Files.exists(absolute)) {
+                    return new SourceFileObject(root);
+                }
+            }
+        }
+        return null;
+    }
+
+    @Override
+    public JavaFileObject getJavaFileForOutput(
+            Location location, String className, JavaFileObject.Kind kind, FileObject sibling) throws IOException {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public FileObject getFileForOutput(Location location, String packageName, String relativeName, FileObject sibling)
+            throws IOException {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void flush() throws IOException {}
+
+    @Override
+    public void close() throws IOException {}
+
+    @Override
+    public int isSupportedOption(String option) {
+        return -1;
+    }
+
+    @Override
+    public Location getLocationForModule(Location location, String moduleName) throws IOException {
+        return delegate.getLocationForModule(location, moduleName);
+    }
+
+    @Override
+    public Location getLocationForModule(Location location, JavaFileObject file) throws IOException {
+        return delegate.getLocationForModule(location, file);
+    }
+
+    @Override
+    public <S> ServiceLoader<S> getServiceLoader(Location location, Class<S> service) throws IOException {
+        getClass().getModule().addUses(service);
+        return ServiceLoader.load(service, getClassLoader(location));
+    }
+
+    @Override
+    public String inferModuleName(Location location) throws IOException {
+        return delegate.inferModuleName(location);
+    }
+
+    @Override
+    public Iterable<Set<Location>> listLocationsForModules(Location location) throws IOException {
+        return delegate.listLocationsForModules(location);
+    }
+
+    @Override
+    public boolean contains(Location location, FileObject file) throws IOException {
+        if (file instanceof SourceFileObject && location == StandardLocation.SOURCE_PATH) {
+            var source = (SourceFileObject) file;
+            return contains(sourcePath, source);
+        } else if (file instanceof JarFileObject) {
+            var jarFile = (JarFileObject) file;
+            for (var jar : classPath) {
+                if (jarFile.jar.equals(jar)) {
+                    return true;
+                }
+            }
+            return false;
+        } else {
+            return delegate.contains(location, file);
+        }
+    }
+
+    private boolean contains(Set<Path> location, SourceFileObject source) {
+        for (var root : location) {
+            if (source.path.startsWith(root)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    @Override
+    public Iterable<? extends JavaFileObject> getJavaFileObjectsFromFiles(Iterable<? extends File> files) {
+        var result = new ArrayList<JavaFileObject>();
+        for (var f : files) {
+            result.add(new SourceFileObject(f.toPath()));
+        }
+        return result;
+    }
+
+    // Just for compatibility with StandardJavaFileManager
+    // TODO delete this once we no longer need the useSourceFileManager flag
+
+    @Override
+    public Iterable<? extends JavaFileObject> getJavaFileObjects(File... files) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public Iterable<? extends JavaFileObject> getJavaFileObjectsFromStrings(Iterable<String> names) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public Iterable<? extends JavaFileObject> getJavaFileObjects(String... names) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void setLocation(Location location, Iterable<? extends File> files) throws IOException {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public Iterable<? extends File> getLocation(Location location) {
+        throw new UnsupportedOperationException();
+    }
+}
diff --git a/src/main/java/org/javacs/SourceFileObject.java b/src/main/java/org/javacs/SourceFileObject.java
new file mode 100644
index 0000000..4310a8a
--- /dev/null
+++ b/src/main/java/org/javacs/SourceFileObject.java
@@ -0,0 +1,105 @@
+package org.javacs;
+
+import java.io.*;
+import java.net.URI;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import javax.lang.model.element.Modifier;
+import javax.lang.model.element.NestingKind;
+import javax.tools.JavaFileObject;
+
+class SourceFileObject implements JavaFileObject {
+    final Path path;
+
+    SourceFileObject(Path path) {
+        this.path = path;
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        if (other.getClass() != SourceFileObject.class) return false;
+        var that = (SourceFileObject) other;
+        return this.path.equals(that.path);
+    }
+
+    @Override
+    public Kind getKind() {
+        var name = path.getFileName().toString();
+        return kindFromExtension(name);
+    }
+
+    static Kind kindFromExtension(String name) {
+        for (var candidate : Kind.values()) {
+            if (name.endsWith(candidate.extension)) {
+                return candidate;
+            }
+        }
+        return null;
+    }
+
+    @Override
+    public boolean isNameCompatible(String simpleName, Kind kind) {
+        var relative = Paths.get(simpleName.replace('.', '/'));
+        return path.endsWith(relative);
+    }
+
+    @Override
+    public NestingKind getNestingKind() {
+        return null;
+    }
+
+    @Override
+    public Modifier getAccessLevel() {
+        return null;
+    }
+
+    @Override
+    public URI toUri() {
+        return path.toUri();
+    }
+
+    @Override
+    public String getName() {
+        return path.toString();
+    }
+
+    @Override
+    public InputStream openInputStream() throws IOException {
+        return Files.newInputStream(path);
+    }
+
+    @Override
+    public OutputStream openOutputStream() throws IOException {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public Reader openReader(boolean ignoreEncodingErrors) throws IOException {
+        return Files.newBufferedReader(path);
+    }
+
+    @Override
+    public CharSequence getCharContent(boolean ignoreEncodingErrors) throws IOException {
+        return Files.readString(path);
+    }
+
+    @Override
+    public Writer openWriter() throws IOException {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public long getLastModified() {
+        try {
+            return Files.getLastModifiedTime(path).toMillis();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public boolean delete() {
+        throw new UnsupportedOperationException();
+    }
+}