Added support for @NonNull in the @Setter annotation

5 files changed, 89 insertions, 8 deletions

diff --git a/src/lombok/eclipse/Eclipse.java b/src/lombok/eclipse/Eclipse.java
index cba2bd05..b1860984 100644
--- a/src/lombok/eclipse/Eclipse.java
+++ b/src/lombok/eclipse/Eclipse.java
@@ -117,7 +117,7 @@ public class Eclipse {
 	 * For 'speed' reasons, Eclipse works a lot with char arrays. I have my doubts this was a fruitful exercise,
 	 * but we need to deal with it. This turns [[java][lang][String]] into "java.lang.String".
 	 */
-	static String toQualifiedName(char[][] typeName) {
+	public static String toQualifiedName(char[][] typeName) {
 		StringBuilder sb = new StringBuilder();
 		boolean first = true;
 		for ( char[] c : typeName ) {
@@ -127,6 +127,16 @@ public class Eclipse {
 		return sb.toString();
 	}
 	
+	public static char[][] fromQualifiedName(String typeName) {
+		String[] split = typeName.split("\\.");
+		char[][] result = new char[split.length][];
+		for (int i = 0; i < split.length; i++) {
+			result[i] = split[i].toCharArray();
+		}
+		return result;
+	}
+	
+	
 	/**
 	 * You can't share TypeParameter objects or bad things happen; for example, one 'T' resolves differently
 	 * from another 'T', even for the same T in a single class file. Unfortunately the TypeParameter type hierarchy
diff --git a/src/lombok/eclipse/handlers/HandleSetter.java b/src/lombok/eclipse/handlers/HandleSetter.java
index adac30d7..567a7fbe 100644
--- a/src/lombok/eclipse/handlers/HandleSetter.java
+++ b/src/lombok/eclipse/handlers/HandleSetter.java
@@ -32,15 +32,24 @@ import lombok.eclipse.EclipseAnnotationHandler;
 import lombok.eclipse.EclipseAST.Node;
 
 import org.eclipse.jdt.internal.compiler.ast.ASTNode;
+import org.eclipse.jdt.internal.compiler.ast.AllocationExpression;
 import org.eclipse.jdt.internal.compiler.ast.Annotation;
 import org.eclipse.jdt.internal.compiler.ast.Argument;
 import org.eclipse.jdt.internal.compiler.ast.Assignment;
+import org.eclipse.jdt.internal.compiler.ast.EqualExpression;
+import org.eclipse.jdt.internal.compiler.ast.Expression;
 import org.eclipse.jdt.internal.compiler.ast.FieldDeclaration;
 import org.eclipse.jdt.internal.compiler.ast.FieldReference;
+import org.eclipse.jdt.internal.compiler.ast.IfStatement;
 import org.eclipse.jdt.internal.compiler.ast.MethodDeclaration;
+import org.eclipse.jdt.internal.compiler.ast.NullLiteral;
+import org.eclipse.jdt.internal.compiler.ast.OperatorIds;
+import org.eclipse.jdt.internal.compiler.ast.QualifiedTypeReference;
 import org.eclipse.jdt.internal.compiler.ast.SingleNameReference;
 import org.eclipse.jdt.internal.compiler.ast.Statement;
+import org.eclipse.jdt.internal.compiler.ast.StringLiteral;
 import org.eclipse.jdt.internal.compiler.ast.ThisReference;
+import org.eclipse.jdt.internal.compiler.ast.ThrowStatement;
 import org.eclipse.jdt.internal.compiler.ast.TypeDeclaration;
 import org.eclipse.jdt.internal.compiler.ast.TypeReference;
 import org.eclipse.jdt.internal.compiler.classfmt.ClassFileConstants;
@@ -119,6 +128,7 @@ public class HandleSetter implements EclipseAnnotationHandler<Setter> {
 	
 	private MethodDeclaration generateSetter(TypeDeclaration parent, FieldDeclaration field, String name,
 			int modifier, ASTNode ast) {
+		
 		long pos = (((long)ast.sourceStart) << 32) | ast.sourceEnd;
 		MethodDeclaration method = new MethodDeclaration(parent.compilationResult);
 		method.modifiers = modifier;
@@ -137,7 +147,22 @@ public class HandleSetter implements EclipseAnnotationHandler<Setter> {
 		Assignment assignment = new Assignment(thisX, new SingleNameReference(field.name, pos), (int)pos);
 		method.bodyStart = method.declarationSourceStart = method.sourceStart = ast.sourceStart;
 		method.bodyEnd = method.declarationSourceEnd = method.sourceEnd = ast.sourceEnd;
-		method.statements = new Statement[] { assignment };
+		
+		Annotation nonNull = findNonNullannotation(field);
+		if (nonNull == null) {
+			method.statements = new Statement[] { assignment };
+		}
+		else {
+			AllocationExpression exception = new AllocationExpression();
+			exception.type = new QualifiedTypeReference(Eclipse.fromQualifiedName("java.lang.NullPointerException"), new long[]{0, 0, 0});
+			exception.arguments = new Expression[] { new StringLiteral(field.name, 0, field.name.length - 1, 0)};
+			ThrowStatement throwStatement = new ThrowStatement(exception, 0, 0);
+			
+			IfStatement nullCheck = new IfStatement(new EqualExpression(new SingleNameReference(field.name, 0),
+					new NullLiteral(0, 0), OperatorIds.EQUAL_EQUAL), throwStatement, 0, 0);
+			
+			method.statements = new Statement[] { nullCheck, assignment };
+		}
 		return method;
 	}
 }
diff --git a/src/lombok/eclipse/handlers/PKG.java b/src/lombok/eclipse/handlers/PKG.java
index 7fdf7afd..db38e3df 100644
--- a/src/lombok/eclipse/handlers/PKG.java
+++ b/src/lombok/eclipse/handlers/PKG.java
@@ -29,9 +29,11 @@ import lombok.eclipse.EclipseAST;
 
 import org.eclipse.jdt.internal.compiler.ast.ASTNode;
 import org.eclipse.jdt.internal.compiler.ast.AbstractMethodDeclaration;
+import org.eclipse.jdt.internal.compiler.ast.Annotation;
 import org.eclipse.jdt.internal.compiler.ast.ConstructorDeclaration;
 import org.eclipse.jdt.internal.compiler.ast.FieldDeclaration;
 import org.eclipse.jdt.internal.compiler.ast.TypeDeclaration;
+import org.eclipse.jdt.internal.compiler.ast.TypeReference;
 
 class PKG {
 	private PKG() {}
@@ -212,4 +214,17 @@ class PKG {
 		
 		type.add(method, Kind.METHOD).recursiveSetHandled();
 	}
+	
+	static Annotation findNonNullannotation(FieldDeclaration field) {
+		for (Annotation annotation : field.annotations) {
+			TypeReference typeRef = annotation.type;
+			if ( typeRef != null && typeRef.getTypeName()!= null ) {
+				char[][] typeName = typeRef.getTypeName();
+				if (new String(typeName[typeName.length - 1]).equals("NonNull")) {
+					return annotation;
+				}
+			}
+		}	
+		return null;
+	}
 }
diff --git a/src/lombok/javac/handlers/HandleSetter.java b/src/lombok/javac/handlers/HandleSetter.java
index d11b4250..b4720387 100644
--- a/src/lombok/javac/handlers/HandleSetter.java
+++ b/src/lombok/javac/handlers/HandleSetter.java
@@ -34,6 +34,8 @@ import lombok.javac.JavacAST.Node;
 import org.mangosdk.spi.ProviderFor;
 
 import com.sun.tools.javac.code.Flags;
+import com.sun.tools.javac.code.TypeTags;
+import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.TreeMaker;
 import com.sun.tools.javac.tree.JCTree.JCAnnotation;
 import com.sun.tools.javac.tree.JCTree.JCAssign;
@@ -119,7 +121,21 @@ public class HandleSetter implements JavacAnnotationHandler<Setter> {
 		JCFieldAccess thisX = treeMaker.Select(treeMaker.Ident(field.toName("this")), fieldDecl.name);
 		JCAssign assign = treeMaker.Assign(thisX, treeMaker.Ident(fieldDecl.name));
 		
-		JCBlock methodBody = treeMaker.Block(0, List.<JCStatement>of(treeMaker.Exec(assign)));
+		
+		List<JCStatement> statements;
+		JCAnnotation nonNull = findNonNullAnnotation(field);
+		if (nonNull == null) {
+			statements = List.<JCStatement>of(treeMaker.Exec(assign));
+		}
+		else {
+			JCExpression npe = chainDots(treeMaker, field, "java", "lang", "NullPointerException");
+			JCTree exception = treeMaker.NewClass(null, List.<JCExpression>nil(), npe, List.<JCExpression>of(treeMaker.Literal(fieldDecl.name.toString())), null);
+			JCStatement throwStatement = treeMaker.Throw(exception);
+			JCStatement nullCheck = treeMaker.If(treeMaker.Binary(JCTree.EQ, treeMaker.Ident(fieldDecl.name), treeMaker.Literal(TypeTags.BOT, null)), throwStatement, null);
+			statements = List.<JCStatement>of(nullCheck, treeMaker.Exec(assign));
+		}
+		
+		JCBlock methodBody = treeMaker.Block(0, statements);
 		Name methodName = field.toName(toSetterName(fieldDecl));
 		JCVariableDecl param = treeMaker.VarDef(treeMaker.Modifiers(0), fieldDecl.name, fieldDecl.vartype, null);
 		JCExpression methodType = treeMaker.Type(field.getSymbolTable().voidType);
diff --git a/src/lombok/javac/handlers/PKG.java b/src/lombok/javac/handlers/PKG.java
index 205a2b6e..bc5b691c 100644
--- a/src/lombok/javac/handlers/PKG.java
+++ b/src/lombok/javac/handlers/PKG.java
@@ -23,20 +23,22 @@ package lombok.javac.handlers;
 
 import java.lang.reflect.Modifier;
 
+import lombok.AccessLevel;
+import lombok.core.TransformationsUtil;
+import lombok.core.AST.Kind;
+import lombok.javac.JavacAST;
+import lombok.javac.JavacAST.Node;
+
 import com.sun.tools.javac.code.Flags;
 import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.TreeMaker;
+import com.sun.tools.javac.tree.JCTree.JCAnnotation;
 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.JCVariableDecl;
 import com.sun.tools.javac.util.List;
 
-import lombok.AccessLevel;
-import lombok.core.TransformationsUtil;
-import lombok.core.AST.Kind;
-import lombok.javac.JavacAST;
-
 /**
  * Container for static utility methods relevant to this package.
  */
@@ -251,4 +253,17 @@ class PKG {
 		
 		return e;
 	}
+	
+	static JCAnnotation findNonNullAnnotation(Node fieldNode) {
+		for ( Node child : fieldNode.down() ) {
+			if ( child.getKind() == Kind.ANNOTATION ) {
+				JCAnnotation annotation = (JCAnnotation) child.get();
+				String name = annotation.annotationType.toString();
+				if (name.equals("NonNull") || name.endsWith(".NonNull")) {
+					return annotation;
+				}
+			}
+		}	
+		return null;
+	}	
 }