Work on HandleData, as well as generalizing features in the the PKG class and updating HandleGetter/Setter to call into it.

6 files changed, 99 insertions, 22 deletions

diff --git a/src/lombok/javac/HandlerLibrary.java b/src/lombok/javac/HandlerLibrary.java
index f9404668..ab5388d3 100644
--- a/src/lombok/javac/HandlerLibrary.java
+++ b/src/lombok/javac/HandlerLibrary.java
@@ -13,6 +13,7 @@ import javax.annotation.processing.Messager;
 import javax.lang.model.element.TypeElement;
 import javax.tools.Diagnostic;
 
+import lombok.core.PrintAST;
 import lombok.core.SpiLoadUtil;
 import lombok.core.TypeLibrary;
 import lombok.core.TypeResolver;
@@ -27,6 +28,7 @@ public class HandlerLibrary {
 	private final Map<String, AnnotationHandlerContainer<?>> annotationHandlers = new HashMap<String, AnnotationHandlerContainer<?>>();
 	private final Collection<JavacASTVisitor> visitorHandlers = new ArrayList<JavacASTVisitor>();
 	private final Messager messager;
+	private boolean skipPrintAST = true;
 	
 	public HandlerLibrary(Messager messager) {
 		this.messager = messager;
@@ -111,6 +113,8 @@ public class HandlerLibrary {
 		String rawType = annotation.annotationType.toString();
 		boolean handled = false;
 		for ( String fqn : resolver.findTypeMatches(node, rawType) ) {
+			boolean isPrintAST = fqn.equals(PrintAST.class.getName());
+			if ( isPrintAST == skipPrintAST ) continue;
 			AnnotationHandlerContainer<?> container = annotationHandlers.get(fqn);
 			if ( container == null ) continue;
 			
@@ -137,4 +141,12 @@ public class HandlerLibrary {
 	public boolean hasHandlerFor(TypeElement annotationType) {
 		return annotationHandlers.containsKey(annotationType.getQualifiedName().toString());
 	}
+	
+	public void skipPrintAST() {
+		skipPrintAST = true;
+	}
+	
+	public void skipAllButPrintAST() {
+		skipPrintAST = false;
+	}
 }
diff --git a/src/lombok/javac/apt/Processor.java b/src/lombok/javac/apt/Processor.java
index 8558a0ba..52b3703a 100644
--- a/src/lombok/javac/apt/Processor.java
+++ b/src/lombok/javac/apt/Processor.java
@@ -56,10 +56,16 @@ public class Processor extends AbstractProcessor {
 		
 		for ( JCCompilationUnit unit : units.keySet() ) asts.add(new JavacAST(trees, processingEnv, unit));
 		
+		handlers.skipPrintAST();
 		for ( JavacAST ast : asts ) {
 			ast.traverse(new AnnotationVisitor());
 			handlers.callASTVisitors(ast);
 		}
+		
+		handlers.skipAllButPrintAST();
+		for ( JavacAST ast : asts ) {
+			ast.traverse(new AnnotationVisitor());
+		}
 		return false;
 	}
 	
diff --git a/src/lombok/javac/handlers/HandleData.java b/src/lombok/javac/handlers/HandleData.java
index 48eb7f76..5d1cc913 100644
--- a/src/lombok/javac/handlers/HandleData.java
+++ b/src/lombok/javac/handlers/HandleData.java
@@ -1,7 +1,8 @@
 package lombok.javac.handlers;
 
-import java.util.ArrayList;
-import java.util.List;
+import java.lang.reflect.Modifier;
+
+import static lombok.javac.handlers.PKG.*;
 
 import lombok.Data;
 import lombok.core.AnnotationValues;
@@ -12,9 +13,17 @@ import lombok.javac.JavacAST.Node;
 import org.mangosdk.spi.ProviderFor;
 
 import com.sun.tools.javac.code.Flags;
+import com.sun.tools.javac.tree.TreeMaker;
 import com.sun.tools.javac.tree.JCTree.JCAnnotation;
+import com.sun.tools.javac.tree.JCTree.JCAssign;
 import com.sun.tools.javac.tree.JCTree.JCClassDecl;
+import com.sun.tools.javac.tree.JCTree.JCExpression;
+import com.sun.tools.javac.tree.JCTree.JCFieldAccess;
+import com.sun.tools.javac.tree.JCTree.JCMethodDecl;
+import com.sun.tools.javac.tree.JCTree.JCModifiers;
+import com.sun.tools.javac.tree.JCTree.JCStatement;
 import com.sun.tools.javac.tree.JCTree.JCVariableDecl;
+import com.sun.tools.javac.util.List;
 
 @ProviderFor(JavacAnnotationHandler.class)
 public class HandleData implements JavacAnnotationHandler<Data> {
@@ -25,25 +34,52 @@ public class HandleData implements JavacAnnotationHandler<Data> {
 		long flags = typeDecl.mods.flags;
 		boolean notAClass = (flags & (Flags.INTERFACE | Flags.ENUM | Flags.ANNOTATION)) != 0;
 		
-		if ( typeDecl != null || notAClass ) {
+		if ( typeDecl == null || notAClass ) {
 			annotationNode.addError("@Data is only supported on a class.");
 			return false;
 		}
 		
-		List<Node> nodesForEquality = new ArrayList<Node>();
+		List<Node> nodesForEquality = List.nil();
+		List<Node> nodesForConstructorAndToString = List.nil();
 		for ( Node child : typeNode.down() ) {
 			if ( child.getKind() != Kind.FIELD ) continue;
 			JCVariableDecl fieldDecl = (JCVariableDecl) child.get();
 			long fieldFlags = fieldDecl.mods.flags;
 			//Skip static fields.
 			if ( (fieldFlags & Flags.STATIC) != 0 ) continue;
-			if ( (fieldFlags & Flags.TRANSIENT) == 0 ) nodesForEquality.add(child);
+			if ( (fieldFlags & Flags.TRANSIENT) == 0 ) nodesForEquality = nodesForEquality.append(child);
+			nodesForConstructorAndToString = nodesForConstructorAndToString.append(child);
 			new HandleGetter().generateGetterForField(child, annotationNode.get());
 			if ( (fieldFlags & Flags.FINAL) == 0 )
 				new HandleSetter().generateSetterForField(child, annotationNode.get());
 		}
 		
-		//TODO generate constructor, hashCode, equals, toString.
+		String staticConstructorName = annotation.getInstance().staticConstructor();
+		JCMethodDecl constructor = createConstructor(staticConstructorName.equals(""), typeNode, nodesForConstructorAndToString);
+		injectMethod(typeNode, constructor);
+		
+		//TODO generate static constructors, hashCode, equals, toString.
 		return true;
 	}
+	
+	private JCMethodDecl createConstructor(boolean isPublic, Node typeNode, List<Node> fields) {
+		TreeMaker maker = typeNode.getTreeMaker();
+		JCClassDecl type = (JCClassDecl) typeNode.get();
+		
+		List<JCStatement> assigns = List.nil();
+		List<JCVariableDecl> params = List.nil();
+		
+		for ( Node fieldNode : fields ) {
+			JCVariableDecl field = (JCVariableDecl) fieldNode.get();
+			JCVariableDecl param = maker.VarDef(maker.Modifiers(0), field.name, field.vartype, null);
+			params = params.append(param);
+			JCFieldAccess thisX = maker.Select(maker.Ident(fieldNode.toName("this")), field.name);
+			JCAssign assign = maker.Assign(thisX, maker.Ident(field.name));
+			assigns = assigns.append(maker.Exec(assign));
+		}
+		
+		JCModifiers mods = maker.Modifiers(isPublic ? Modifier.PUBLIC : Modifier.PRIVATE);
+		return maker.MethodDef(mods,
+				typeNode.toName("<init>"), null, type.typarams, params, List.<JCExpression>nil(), maker.Block(0L, assigns), null);
+	}
 }
diff --git a/src/lombok/javac/handlers/HandleGetter.java b/src/lombok/javac/handlers/HandleGetter.java
index 3405810b..f66247ff 100644
--- a/src/lombok/javac/handlers/HandleGetter.java
+++ b/src/lombok/javac/handlers/HandleGetter.java
@@ -1,7 +1,6 @@
 package lombok.javac.handlers;
 
 import static lombok.javac.handlers.PKG.*;
-
 import lombok.AccessLevel;
 import lombok.Getter;
 import lombok.core.AnnotationValues;
@@ -16,7 +15,6 @@ import com.sun.tools.javac.code.Flags;
 import com.sun.tools.javac.tree.TreeMaker;
 import com.sun.tools.javac.tree.JCTree.JCAnnotation;
 import com.sun.tools.javac.tree.JCTree.JCBlock;
-import com.sun.tools.javac.tree.JCTree.JCClassDecl;
 import com.sun.tools.javac.tree.JCTree.JCExpression;
 import com.sun.tools.javac.tree.JCTree.JCMethodDecl;
 import com.sun.tools.javac.tree.JCTree.JCStatement;
@@ -75,14 +73,9 @@ public class HandleGetter implements JavacAnnotationHandler<Getter> {
 			//continue with creating the getter
 		}
 		
-		JCClassDecl javacClassTree = (JCClassDecl) fieldNode.up().get();
-		
 		long access = toJavacModifier(level) | (fieldDecl.mods.flags & Flags.STATIC);
 		
-		JCMethodDecl getterMethod = createGetter(access, fieldNode, fieldNode.getTreeMaker());
-		javacClassTree.defs = javacClassTree.defs.append(getterMethod);
-		
-		fieldNode.up().add(getterMethod, Kind.METHOD).recursiveSetHandled();
+		injectMethod(fieldNode.up(), createGetter(access, fieldNode, fieldNode.getTreeMaker()));
 		
 		return true;
 	}
diff --git a/src/lombok/javac/handlers/HandleSetter.java b/src/lombok/javac/handlers/HandleSetter.java
index 4999b5d8..2b766db1 100644
--- a/src/lombok/javac/handlers/HandleSetter.java
+++ b/src/lombok/javac/handlers/HandleSetter.java
@@ -17,7 +17,6 @@ import com.sun.tools.javac.tree.TreeMaker;
 import com.sun.tools.javac.tree.JCTree.JCAnnotation;
 import com.sun.tools.javac.tree.JCTree.JCAssign;
 import com.sun.tools.javac.tree.JCTree.JCBlock;
-import com.sun.tools.javac.tree.JCTree.JCClassDecl;
 import com.sun.tools.javac.tree.JCTree.JCExpression;
 import com.sun.tools.javac.tree.JCTree.JCFieldAccess;
 import com.sun.tools.javac.tree.JCTree.JCMethodDecl;
@@ -78,14 +77,9 @@ public class HandleSetter implements JavacAnnotationHandler<Setter> {
 			//continue with creating the setter
 		}
 		
-		JCClassDecl javacClassTree = (JCClassDecl) fieldNode.up().get();
-		
 		long access = toJavacModifier(level) | (fieldDecl.mods.flags & Flags.STATIC);
 		
-		JCMethodDecl setterMethod = createSetter(access, fieldNode, fieldNode.getTreeMaker());
-		javacClassTree.defs = javacClassTree.defs.append(setterMethod);
-		
-		fieldNode.up().add(setterMethod, Kind.METHOD).recursiveSetHandled();
+		injectMethod(fieldNode.up(), createSetter(access, fieldNode, fieldNode.getTreeMaker()));
 		
 		return true;
 	}
@@ -98,7 +92,7 @@ public class HandleSetter implements JavacAnnotationHandler<Setter> {
 		
 		JCBlock methodBody = treeMaker.Block(0, List.<JCStatement>of(treeMaker.Exec(assign)));
 		Name methodName = field.toName(toSetterName(fieldDecl));
-		JCVariableDecl param = treeMaker.VarDef(treeMaker.Modifiers(0, List.<JCAnnotation>nil()), fieldDecl.name, fieldDecl.vartype, null);
+		JCVariableDecl param = treeMaker.VarDef(treeMaker.Modifiers(0), fieldDecl.name, fieldDecl.vartype, null);
 		JCExpression methodType = treeMaker.Type(field.getSymbolTable().voidType);
 		
 		List<JCTypeParameter> methodGenericParams = List.nil();
diff --git a/src/lombok/javac/handlers/PKG.java b/src/lombok/javac/handlers/PKG.java
index 138559a9..615e4c80 100644
--- a/src/lombok/javac/handlers/PKG.java
+++ b/src/lombok/javac/handlers/PKG.java
@@ -2,13 +2,16 @@ package lombok.javac.handlers;
 
 import java.lang.reflect.Modifier;
 
+import com.sun.tools.javac.code.Flags;
 import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.JCTree.JCClassDecl;
 import com.sun.tools.javac.tree.JCTree.JCMethodDecl;
 import com.sun.tools.javac.tree.JCTree.JCVariableDecl;
+import com.sun.tools.javac.util.List;
 
 import lombok.AccessLevel;
 import lombok.core.TransformationsUtil;
+import lombok.core.AST.Kind;
 import lombok.javac.JavacAST;
 
 class PKG {
@@ -66,4 +69,37 @@ class PKG {
 			return Modifier.PROTECTED;
 		}
 	}
+	
+	static void injectMethod(JavacAST.Node typeNode, JCMethodDecl method) {
+		JCClassDecl type = (JCClassDecl) typeNode.get();
+		
+		if ( method.getName().contentEquals("<init>") ) {
+			//Scan for default constructor, and remove it.
+			int idx = 0;
+			for ( JCTree def : type.defs ) {
+				if ( def instanceof JCMethodDecl ) {
+					if ( (((JCMethodDecl)def).mods.flags & Flags.GENERATEDCONSTR) != 0 ) {
+						JavacAST.Node tossMe = typeNode.getNodeFor(def);
+						if ( tossMe != null ) tossMe.up().removeChild(tossMe);
+						type.defs = addAllButOne(type.defs, idx);
+						break;
+					}
+				}
+				idx++;
+			}
+		}
+		
+		type.defs = type.defs.append(method);
+		
+		typeNode.add(method, Kind.METHOD).recursiveSetHandled();
+	}
+	
+	private static List<JCTree> addAllButOne(List<JCTree> defs, int idx) {
+		List<JCTree> out = List.nil();
+		int i = 0;
+		for ( JCTree def : defs ) {
+			if ( i++ != idx ) out = out.append(def);
+		}
+		return out;
+	}
 }