Post-compilation works in both Eclipse and ecj

2 files changed, 47 insertions, 0 deletions

diff --git a/src/eclipseAgent/lombok/eclipse/agent/EclipsePatcher.java b/src/eclipseAgent/lombok/eclipse/agent/EclipsePatcher.java
index 1486c608..1d2efea0 100644
--- a/src/eclipseAgent/lombok/eclipse/agent/EclipsePatcher.java
+++ b/src/eclipseAgent/lombok/eclipse/agent/EclipsePatcher.java
@@ -64,11 +64,39 @@ public class EclipsePatcher extends Agent {
 			patchSetGeneratedFlag(sm);
 			patchHideGeneratedNodes(sm);
 			patchLiveDebug(sm);
+			patchPostCompileHookEclipse(sm);
+		} else {
+			patchPostCompileHookEcj(sm);
 		}
 		
 		if (reloadExistingClasses) sm.reloadClasses(instrumentation);
 	}
 	
+	private static void patchPostCompileHookEclipse(ScriptManager sm) {
+		sm.addScript(ScriptBuilder.wrapMethodCall()
+				.target(new MethodTarget("org.eclipse.jdt.internal.core.builder.IncrementalImageBuilder", "writeClassFileContents"))
+				.target(new MethodTarget("org.eclipse.jdt.internal.core.builder.AbstractImageBuilder", "writeClassFileContents"))
+				.methodToWrap(new Hook("org.eclipse.jdt.internal.compiler.ClassFile", "getBytes", "byte[]"))
+				.wrapMethod(new Hook("lombok.eclipse.agent.PatchFixes", "runPostCompiler", "byte[]", "byte[]", "java.lang.String"))
+				.requestExtra(StackRequest.PARAM3)
+				.transplant()
+				.build());
+	}
+	
+	private static void patchPostCompileHookEcj(ScriptManager sm) {
+		sm.addScript(ScriptBuilder.wrapMethodCall()
+				.target(new MethodTarget("org.eclipse.jdt.internal.compiler.tool.EclipseCompilerImpl", "outputClassFiles"))
+				.methodToWrap(new Hook("javax.tools.JavaFileObject", "openOutputStream", "java.io.OutputStream"))
+				.wrapMethod(new Hook("lombok.eclipse.agent.PatchFixes", "runPostCompiler", "java.io.OutputStream", "java.io.OutputStream"))
+				.build());
+		
+		sm.addScript(ScriptBuilder.wrapMethodCall()
+				.target(new MethodTarget("org.eclipse.jdt.internal.compiler.util.Util", "writeToDisk"))
+				.methodToWrap(new Hook("java.io.BufferedOutputStream", "<init>", "void", "java.io.OutputStream", "int"))
+				.wrapMethod(new Hook("lombok.eclipse.agent.PatchFixes", "runPostCompiler", "java.io.BufferedOutputStream", "java.io.BufferedOutputStream"))
+				.build());
+	}
+	
 	private static void patchHideGeneratedNodes(ScriptManager sm) {
 		sm.addScript(ScriptBuilder.wrapReturnValue()
 				.target(new MethodTarget("org.eclipse.jdt.internal.corext.dom.LinkedNodeFinder", "findByNode"))
diff --git a/src/eclipseAgent/lombok/eclipse/agent/PatchFixes.java b/src/eclipseAgent/lombok/eclipse/agent/PatchFixes.java
index bb3d90f7..48103a98 100644
--- a/src/eclipseAgent/lombok/eclipse/agent/PatchFixes.java
+++ b/src/eclipseAgent/lombok/eclipse/agent/PatchFixes.java
@@ -22,10 +22,16 @@
 
 package lombok.eclipse.agent;
 
+import java.io.BufferedOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
 import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.List;
 
+import lombok.core.DiagnosticsReceiver;
+import lombok.core.PostCompiler;
+
 import org.eclipse.jdt.core.IMethod;
 import org.eclipse.jdt.core.dom.ASTNode;
 import org.eclipse.jdt.core.dom.SimpleName;
@@ -107,4 +113,17 @@ public class PatchFixes {
 		}
 		return newSimpleNames;
 	}
+	
+	public static byte[] runPostCompiler(byte[] bytes, String className) {
+		byte[] transformed = PostCompiler.applyTransformations(bytes, className, DiagnosticsReceiver.CONSOLE);
+		return transformed == null ? bytes : transformed;
+	}
+	
+	public static OutputStream runPostCompiler(OutputStream out) throws IOException {
+		return PostCompiler.wrapOutputStream(out, "TEST", DiagnosticsReceiver.CONSOLE);
+	}
+	
+	public static BufferedOutputStream runPostCompiler(BufferedOutputStream out) throws IOException {
+		return new BufferedOutputStream(PostCompiler.wrapOutputStream(out, "TEST", DiagnosticsReceiver.CONSOLE));
+	}
 }