[i2078] Add possibility to generate assert on `@NonNull`

2 files changed, 48 insertions, 33 deletions

diff --git a/src/core/lombok/javac/handlers/HandleNonNull.java b/src/core/lombok/javac/handlers/HandleNonNull.java
index 9a81ffff..49b987ce 100644
--- a/src/core/lombok/javac/handlers/HandleNonNull.java
+++ b/src/core/lombok/javac/handlers/HandleNonNull.java
@@ -21,13 +21,16 @@
  */
 package lombok.javac.handlers;
 
-import static lombok.core.handlers.HandlerUtil.*;
+import static lombok.core.handlers.HandlerUtil.handleFlagUsage;
 import static lombok.javac.Javac.*;
+import static lombok.javac.JavacTreeMaker.TreeTag.treeTag;
+import static lombok.javac.JavacTreeMaker.TypeTag.typeTag;
 import static lombok.javac.handlers.JavacHandlerUtil.*;
 
 import org.mangosdk.spi.ProviderFor;
 
 import com.sun.tools.javac.tree.JCTree.JCAnnotation;
+import com.sun.tools.javac.tree.JCTree.JCAssert;
 import com.sun.tools.javac.tree.JCTree.JCBinary;
 import com.sun.tools.javac.tree.JCTree.JCBlock;
 import com.sun.tools.javac.tree.JCTree.JCExpression;
@@ -45,13 +48,11 @@ import com.sun.tools.javac.util.List;
 
 import lombok.ConfigurationKeys;
 import lombok.NonNull;
+import lombok.core.AST.Kind;
 import lombok.core.AnnotationValues;
 import lombok.core.HandlerPriority;
-import lombok.core.AST.Kind;
 import lombok.javac.JavacAnnotationHandler;
 import lombok.javac.JavacNode;
-import static lombok.javac.JavacTreeMaker.TypeTag.*;
-import static lombok.javac.JavacTreeMaker.TreeTag.*;
 
 @ProviderFor(JavacAnnotationHandler.class)
 @HandlerPriority(value = 512) // 2^9; onParameter=@__(@NonNull) has to run first.
@@ -161,14 +162,16 @@ public class HandleNonNull extends JavacAnnotationHandler<NonNull> {
 	}
 	
 	/**
-	 * Checks if the statement is of the form 'if (x == null) {throw WHATEVER;},
+	 * Checks if the statement is of the form 'if (x == null) {throw WHATEVER;}' or 'assert x != null;',
 	 * where the block braces are optional. If it is of this form, returns "x".
 	 * If it is not of this form, returns null.
 	 */
 	public String returnVarNameIfNullCheck(JCStatement stat) {
-		if (!(stat instanceof JCIf)) return null;
+		boolean isIf = stat instanceof JCIf; 
+		if (!isIf && !(stat instanceof JCAssert)) return null;
 		
-		/* Check that the if's statement is a throw statement, possibly in a block. */ {
+		if (isIf) {
+			/* Check that the if's statement is a throw statement, possibly in a block. */
 			JCStatement then = ((JCIf) stat).thenpart;
 			if (then instanceof JCBlock) {
 				List<JCStatement> stats = ((JCBlock) then).stats;
@@ -180,11 +183,15 @@ public class HandleNonNull extends JavacAnnotationHandler<NonNull> {
 		
 		/* Check that the if's conditional is like 'x == null'. Return from this method (don't generate
 		   a nullcheck) if 'x' is equal to our own variable's name: There's already a nullcheck here. */ {
-			JCExpression cond = ((JCIf) stat).cond;
+			JCExpression cond = isIf ? ((JCIf) stat).cond : ((JCAssert) stat).cond;
 			while (cond instanceof JCParens) cond = ((JCParens) cond).expr;
 			if (!(cond instanceof JCBinary)) return null;
 			JCBinary bin = (JCBinary) cond;
-			if (!CTC_EQUAL.equals(treeTag(bin))) return null;
+			if (isIf) {
+				if (!CTC_EQUAL.equals(treeTag(bin))) return null;
+			} else {
+				if (!CTC_NOT_EQUAL.equals(treeTag(bin))) return null;
+			}
 			if (!(bin.lhs instanceof JCIdent)) return null;
 			if (!(bin.rhs instanceof JCLiteral)) return null;
 			if (!CTC_BOT.equals(typeTag(bin.rhs))) return null;
diff --git a/src/core/lombok/javac/handlers/JavacHandlerUtil.java b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
index 9e6c8c00..509a7397 100644
--- a/src/core/lombok/javac/handlers/JavacHandlerUtil.java
+++ b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
@@ -40,36 +40,15 @@ import java.util.regex.Pattern;
 
 import javax.lang.model.element.Element;
 
-import lombok.AccessLevel;
-import lombok.ConfigurationKeys;
-import lombok.Data;
-import lombok.Getter;
-import lombok.core.AST.Kind;
-import lombok.core.AnnotationValues;
-import lombok.core.LombokImmutableList;
-import lombok.core.AnnotationValues.AnnotationValue;
-import lombok.core.CleanupTask;
-import lombok.core.TypeResolver;
-import lombok.core.configuration.NullCheckExceptionType;
-import lombok.core.configuration.TypeName;
-import lombok.core.handlers.HandlerUtil;
-import lombok.delombok.LombokOptionsFactory;
-import lombok.experimental.Accessors;
-import lombok.experimental.Tolerate;
-import lombok.javac.Javac;
-import lombok.javac.JavacNode;
-import lombok.javac.JavacTreeMaker;
-import lombok.permit.Permit;
-
 import com.sun.tools.javac.code.BoundKind;
 import com.sun.tools.javac.code.Flags;
 import com.sun.tools.javac.code.Scope;
 import com.sun.tools.javac.code.Symbol;
-import com.sun.tools.javac.code.Symtab;
-import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.Symbol.ClassSymbol;
 import com.sun.tools.javac.code.Symbol.MethodSymbol;
 import com.sun.tools.javac.code.Symbol.VarSymbol;
+import com.sun.tools.javac.code.Symtab;
+import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.Type.MethodType;
 import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.JCTree.JCAnnotation;
@@ -83,6 +62,7 @@ import com.sun.tools.javac.tree.JCTree.JCExpressionStatement;
 import com.sun.tools.javac.tree.JCTree.JCFieldAccess;
 import com.sun.tools.javac.tree.JCTree.JCIdent;
 import com.sun.tools.javac.tree.JCTree.JCImport;
+import com.sun.tools.javac.tree.JCTree.JCLiteral;
 import com.sun.tools.javac.tree.JCTree.JCMethodDecl;
 import com.sun.tools.javac.tree.JCTree.JCMethodInvocation;
 import com.sun.tools.javac.tree.JCTree.JCModifiers;
@@ -103,6 +83,28 @@ import com.sun.tools.javac.util.ListBuffer;
 import com.sun.tools.javac.util.Name;
 import com.sun.tools.javac.util.Options;
 
+import lombok.AccessLevel;
+import lombok.ConfigurationKeys;
+import lombok.Data;
+import lombok.Getter;
+import lombok.core.AST.Kind;
+import lombok.core.AnnotationValues;
+import lombok.core.AnnotationValues.AnnotationValue;
+import lombok.core.CleanupTask;
+import lombok.core.LombokImmutableList;
+import lombok.core.TypeResolver;
+import lombok.core.configuration.NullCheckExceptionType;
+import lombok.core.configuration.TypeName;
+import lombok.core.handlers.HandlerUtil;
+import lombok.core.handlers.HandlerUtil.FieldAccess;
+import lombok.delombok.LombokOptionsFactory;
+import lombok.experimental.Accessors;
+import lombok.experimental.Tolerate;
+import lombok.javac.Javac;
+import lombok.javac.JavacNode;
+import lombok.javac.JavacTreeMaker;
+import lombok.permit.Permit;
+
 /**
  * Container for static utility methods useful to handlers written for javac.
  */
@@ -1484,8 +1486,14 @@ public class JavacHandlerUtil {
 		
 		if (isPrimitive(varDecl.vartype)) return null;
 		Name fieldName = varDecl.name;
+		
+		JCLiteral message = maker.Literal(exceptionType.toExceptionMessage(fieldName.toString()));
+		if (exceptionType == NullCheckExceptionType.ASSERTION) {
+			return maker.Assert(maker.Binary(CTC_NOT_EQUAL, maker.Ident(fieldName), maker.Literal(CTC_BOT, null)), message);
+		}
+		
 		JCExpression exType = genTypeRef(variable, exceptionType.getExceptionType());
-		JCExpression exception = maker.NewClass(null, List.<JCExpression>nil(), exType, List.<JCExpression>of(maker.Literal(exceptionType.toExceptionMessage(fieldName.toString()))), null);
+		JCExpression exception = maker.NewClass(null, List.<JCExpression>nil(), exType, List.<JCExpression>of(message), null);
 		JCStatement throwStatement = maker.Throw(exception);
 		JCBlock throwBlock = maker.Block(0, List.of(throwStatement));
 		return maker.If(maker.Binary(CTC_EQUAL, maker.Ident(fieldName), maker.Literal(CTC_BOT, null)), throwBlock, null);