mem cls file manager improve

Author: pfmiles

src/main/java/com/github/pfmiles/dropincc/impl/hotcompile/HotCompileUtil.java

@@ -51,8 +51,10 @@ public static CompilationResult compile(String qualifiedName, String sourceCode)
 
         JavaFileManager fileManager = null;
         Map<String, JavaMemCls> clses = new HashMap<String, JavaMemCls>();
+        Map<String, JavaStringSource> srcs = new HashMap<String, JavaStringSource>();
+        srcs.put(source.getClsName(), source);
         try {
-            fileManager = new MemClsFileManager(compiler.getStandardFileManager(null, null, null), clses);
+            fileManager = new MemClsFileManager(compiler.getStandardFileManager(null, null, null), clses, srcs);
             DiagnosticCollector<JavaFileObject> diagnostics = new DiagnosticCollector<JavaFileObject>();
             StringWriter out = new StringWriter();
             CompilationTask task = compiler.getTask(out, fileManager, diagnostics, options, null, ss);

src/main/java/com/github/pfmiles/dropincc/impl/hotcompile/JavaMemCls.java

@@ -15,6 +15,7 @@
  */
 public class JavaMemCls extends SimpleJavaFileObject {
 
+    private String clsName;
     private ByteArrayOutputStream bos;
 
     protected JavaMemCls(String name) {
@@ -29,4 +30,8 @@ public byte[] getClsBytes() {
     public OutputStream openOutputStream() throws IOException {
         return bos;
     }
+
+    public String getClsName() {
+        return clsName;
+    }
 }

src/main/java/com/github/pfmiles/dropincc/impl/hotcompile/JavaStringSource.java

@@ -21,22 +21,22 @@
  */
 public class JavaStringSource extends SimpleJavaFileObject {
 
-    private String name;
+    private String clsName;
     // source code
     private String source;
 
     protected JavaStringSource(String name, String source) {
         super(URI.create("string:///" + name.replaceAll("\\.", "/") + Kind.SOURCE.extension), Kind.SOURCE);
-        this.name = name;
+        this.clsName = name;
         this.source = source;
     }
 
     public CharSequence getCharContent(boolean ignoreEncodingErrors) throws IOException {
         return this.source;
     }
 
-    public String getName() {
-        return name;
+    public String getClsName() {
+        return clsName;
     }
 
     public String getSource() {

src/main/java/com/github/pfmiles/dropincc/impl/hotcompile/MemClsFileManager.java

@@ -1,32 +1,57 @@
 package com.github.pfmiles.dropincc.impl.hotcompile;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import javax.tools.FileObject;
 import javax.tools.ForwardingJavaFileManager;
 import javax.tools.JavaFileObject;
 import javax.tools.JavaFileObject.Kind;
 import javax.tools.StandardJavaFileManager;
+import javax.tools.StandardLocation;
 
 /**
  * A forwarding file manager which is used to decouple the file system
  * dependency. That is, the compilation process write compiled class files into
  * memory instead of file system.
  * 
+ * <p>
+ * These following methods would be invoked during the compilation progress, so
+ * they should behave well for sure:
+ * </p>
+ * 
+ * <pre>
+ * close
+ * flush
+ * getClassLoader
+ * getJavaFileForOutput
+ * handleOption
+ * hasLocation
+ * inferBinaryName
+ * list
+ * </pre>
+ * 
  * @author pf-miles
  * 
  */
 public class MemClsFileManager extends ForwardingJavaFileManager<StandardJavaFileManager> {
 
     private Map<String, JavaMemCls> destFiles;
+    private Map<String, JavaStringSource> srcFiles;
 
-    protected MemClsFileManager(StandardJavaFileManager fileManager, Map<String, JavaMemCls> destFiles) {
+    protected MemClsFileManager(StandardJavaFileManager fileManager, Map<String, JavaMemCls> destFiles, Map<String, JavaStringSource> srcFiles) {
         super(fileManager);
         this.destFiles = destFiles;
+        this.srcFiles = srcFiles;
     }
 
+    // redirects class file output to memory
     public JavaFileObject getJavaFileForOutput(Location location, String className, Kind kind, FileObject sibling) throws IOException {
+        if (!(Kind.CLASS.equals(kind) && StandardLocation.CLASS_OUTPUT.equals(location)))
+            return super.getJavaFileForOutput(location, className, kind, sibling);
         if (destFiles.containsKey(className)) {
             return destFiles.get(className);
         } else {
@@ -41,4 +66,51 @@ public void close() throws IOException {
         this.destFiles = null;
     }
 
+    public Iterable<JavaFileObject> list(Location location, String packageName, Set<Kind> kinds, boolean recurse) throws IOException {
+        List<JavaFileObject> ret = new ArrayList<JavaFileObject>();
+        if ((StandardLocation.CLASS_OUTPUT.equals(location) || StandardLocation.CLASS_PATH.equals(location)) && kinds.contains(Kind.CLASS)) {
+            for (Map.Entry<String, JavaMemCls> e : destFiles.entrySet()) {
+                String pkgName = resolvePkgName(e.getKey());
+                if (recurse) {
+                    if (pkgName.contains(packageName))
+                        ret.add(e.getValue());
+                } else {
+                    if (pkgName.equals(packageName))
+                        ret.add(e.getValue());
+                }
+            }
+        } else if (StandardLocation.SOURCE_PATH.equals(location) && kinds.contains(Kind.SOURCE)) {
+            for (Map.Entry<String, JavaStringSource> e : srcFiles.entrySet()) {
+                String pkgName = resolvePkgName(e.getKey());
+                if (recurse) {
+                    if (pkgName.contains(packageName))
+                        ret.add(e.getValue());
+                } else {
+                    if (pkgName.equals(packageName))
+                        ret.add(e.getValue());
+                }
+            }
+        }
+        // 也包含super.list
+        Iterable<JavaFileObject> superList = super.list(location, packageName, kinds, recurse);
+        if (superList != null)
+            for (JavaFileObject f : superList)
+                ret.add(f);
+        return ret;
+    }
+
+    private String resolvePkgName(String fullQualifiedClsName) {
+        return fullQualifiedClsName.substring(0, fullQualifiedClsName.lastIndexOf('.'));
+    }
+
+    public String inferBinaryName(Location location, JavaFileObject file) {
+        if (file instanceof JavaMemCls) {
+            return ((JavaMemCls) file).getClsName();
+        } else if (file instanceof JavaStringSource) {
+            return ((JavaStringSource) file).getClsName();
+        } else {
+            return super.inferBinaryName(location, file);
+        }
+    }
+
 }