Removed the WeakHashMap for caching EclipseAST objects; obviously wasn't working due to

circular reference from the EclipseAST back to the CUD.

Now, patched a field into CompilationUnitDeclaration and using that, which works much better
together with the garbage collector.

3 files changed, 51 insertions, 4 deletions

diff --git a/src_eclipseagent/lombok/agent/eclipse/EclipseCUDTransformer.java b/src_eclipseagent/lombok/agent/eclipse/EclipseCUDTransformer.java
new file mode 100644
index 00000000..a8d32024
--- /dev/null
+++ b/src_eclipseagent/lombok/agent/eclipse/EclipseCUDTransformer.java
@@ -0,0 +1,37 @@
+package lombok.agent.eclipse;
+
+import org.objectweb.asm.ClassAdapter;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.FieldVisitor;
+import org.objectweb.asm.Opcodes;
+
+class EclipseCUDTransformer {
+	private final byte[] in;
+
+	EclipseCUDTransformer(byte[] classfileBuffer) {
+		in = classfileBuffer;
+	}
+
+	byte[] transform() {
+		ClassReader reader = new ClassReader(in);
+		ClassWriter writer = new ClassWriter(reader, 0);
+	
+		ClassAdapter adapter = new CUDPatcherAdapter(writer);
+		reader.accept(adapter, 0);
+		return writer.toByteArray();
+	}
+	
+	private static class CUDPatcherAdapter extends ClassAdapter {
+		CUDPatcherAdapter(ClassVisitor cv) {
+			super(cv);
+		}
+		
+		@Override public void visitEnd() {
+			FieldVisitor fv = cv.visitField(Opcodes.ACC_PUBLIC, "$lombokAST", "Ljava/lang/Object;", null, null);
+			fv.visitEnd();
+			cv.visitEnd();
+		}
+	}
+}
diff --git a/src_eclipseagent/lombok/agent/eclipse/EclipseParserPatcher.java b/src_eclipseagent/lombok/agent/eclipse/EclipseParserPatcher.java
index abe26683..81309b34 100644
--- a/src_eclipseagent/lombok/agent/eclipse/EclipseParserPatcher.java
+++ b/src_eclipseagent/lombok/agent/eclipse/EclipseParserPatcher.java
@@ -13,12 +13,21 @@ public class EclipseParserPatcher {
 				ProtectionDomain protectionDomain, byte[] classfileBuffer)
 				throws IllegalClassFormatException {
 			
-			if ( !ECLIPSE_PARSER_CLASS_NAME.equals(className) ) return null;
-			EclipseParserTransformer transformer = new EclipseParserTransformer(classfileBuffer);
-			return transformer.transform();
+			if ( ECLIPSE_PARSER_CLASS_NAME.equals(className) ) {
+				EclipseParserTransformer transformer = new EclipseParserTransformer(classfileBuffer);
+				return transformer.transform();
+			}
+			
+			if ( ECLIPSE_CUD_CLASS_NAME.equals(className) ) {
+				EclipseCUDTransformer transformer = new EclipseCUDTransformer(classfileBuffer);
+				return transformer.transform();
+			}
+			
+			return null;
 		}
 	}
-
+	
+	static final String ECLIPSE_CUD_CLASS_NAME = "org/eclipse/jdt/internal/compiler/ast/CompilationUnitDeclaration";
 	static final String ECLIPSE_PARSER_CLASS_NAME = "org/eclipse/jdt/internal/compiler/parser/Parser";
 	
 	public static void agentmain(String agentArgs, Instrumentation instrumentation) {
diff --git a/src_eclipseagent/lombok/agent/eclipse/EclipseParserTransformer.java b/src_eclipseagent/lombok/agent/eclipse/EclipseParserTransformer.java
index 4a1b4a01..970aafec 100644
--- a/src_eclipseagent/lombok/agent/eclipse/EclipseParserTransformer.java
+++ b/src_eclipseagent/lombok/agent/eclipse/EclipseParserTransformer.java
@@ -43,6 +43,7 @@ class EclipseParserTransformer {
 	byte[] transform() {
 		ClassReader reader = new ClassReader(in);
 		ClassWriter writer = new ClassWriter(reader, 0);
+		
 		ClassAdapter adapter = new ParserPatcherAdapter(writer);
 		reader.accept(adapter, 0);
 		return writer.toByteArray();