Fixed issue with inner classes

src/main/java/org/mdkt/compiler/ExtendedStandardJavaFileManager.java

@@ -2,20 +2,23 @@
 
 import java.io.FileNotFoundException;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
 
 import javax.tools.FileObject;
 import javax.tools.ForwardingJavaFileManager;
 import javax.tools.JavaFileManager;
 import javax.tools.JavaFileObject;
 
 /**
- * Created by trung on 5/3/15.
- * Edited by turpid-monkey on 9/25/15, completed support for multiple compile units.
+ * Created by trung on 5/3/15. Edited by turpid-monkey on 9/25/15, completed
+ * support for multiple compile units.
  */
 public class ExtendedStandardJavaFileManager extends
 		ForwardingJavaFileManager<JavaFileManager> {
 
-	private CompiledCode[] compiledCode;
+	private List<CompiledCode> compiledCode = new ArrayList<CompiledCode>();
 	private DynamicClassLoader cl;
 
 	/**
@@ -26,23 +29,27 @@ public class ExtendedStandardJavaFileManager extends
 	 * @param cl
 	 */
 	protected ExtendedStandardJavaFileManager(JavaFileManager fileManager,
-			DynamicClassLoader cl, CompiledCode... compiledCode) {
+			DynamicClassLoader cl) {
 		super(fileManager);
-		this.compiledCode = compiledCode;
 		this.cl = cl;
-		for (CompiledCode code : compiledCode) {
-			this.cl.addCode(code);
-		}
 	}
 
 	@Override
-    public JavaFileObject getJavaFileForOutput(JavaFileManager.Location location, String className, JavaFileObject.Kind kind, FileObject sibling) throws IOException {
-    	for (CompiledCode code : compiledCode)
-    	{
-    		if (code.getClassName().equals(className)) return code;
-    	}
-    	throw new FileNotFoundException("Missing source code for class " + className );
-    }
+	public JavaFileObject getJavaFileForOutput(
+			JavaFileManager.Location location, String className,
+			JavaFileObject.Kind kind, FileObject sibling) throws IOException {
+
+		try {
+			CompiledCode innerClass = new CompiledCode(className);
+			compiledCode.add(innerClass);
+			cl.addCode(innerClass);
+			return innerClass;
+		} catch (Exception e) {
+			throw new RuntimeException(
+					"Error while creating in-memory output file for "
+							+ className, e);
+		}
+	}
 
 	@Override
 	public ClassLoader getClassLoader(JavaFileManager.Location location) {

src/main/java/org/mdkt/compiler/InMemoryJavaCompiler.java

@@ -1,11 +1,15 @@
 package org.mdkt.compiler;
 
+import java.io.IOException;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
 
+import javax.tools.Diagnostic;
+import javax.tools.DiagnosticListener;
 import javax.tools.JavaCompiler;
+import javax.tools.JavaFileObject;
 import javax.tools.ToolProvider;
 
 /**
@@ -49,7 +53,7 @@ public Map<String, Class<?>> compileAll() throws Exception {
 		}
 
 		ExtendedStandardJavaFileManager fileManager = new ExtendedStandardJavaFileManager(
-				javac.getStandardFileManager(null, null, null), classLoader, code);
+				javac.getStandardFileManager(null, null, null), classLoader);
 		JavaCompiler.CompilationTask task = javac.getTask(null, fileManager,
 				null, null, null, compilationUnits);
 		boolean result = task.call();

src/test/java/org/mdkt/compiler/InMemoryJavaCompilerTest.java

@@ -42,4 +42,19 @@ public void compile_severalFiles() throws Exception {
         Object a = aClass.newInstance();
         Assert.assertEquals("B!", aClass.getMethod("b").invoke(a).toString());
     }
+
+    @Test
+    public void compile_filesWithInnerClasses() throws Exception {
+    	StringBuffer sourceCode = new StringBuffer();
+
+        sourceCode.append("package org.mdkt;\n");
+        sourceCode.append("public class HelloClass {\n");
+        sourceCode.append("   private static class InnerHelloWorld { int inner; }\n");
+        sourceCode.append("   public String hello() { return \"hello\"; }");
+        sourceCode.append("}");
+
+        Class<?> helloClass = InMemoryJavaCompiler.compile("org.mdkt.HelloClass", sourceCode.toString());
+        Assert.assertNotNull(helloClass);
+        Assert.assertEquals(1, helloClass.getDeclaredMethods().length);
+    }
 }