[records] [javac] Putting `@NonNull` on a 'field' of a record now affects the compact constructors, generating one if need be.

2 files changed, 141 insertions, 74 deletions

diff --git a/src/core/lombok/javac/handlers/HandleNonNull.java b/src/core/lombok/javac/handlers/HandleNonNull.java
index 2f359c0f..8e5b0030 100644
--- a/src/core/lombok/javac/handlers/HandleNonNull.java
+++ b/src/core/lombok/javac/handlers/HandleNonNull.java
@@ -27,11 +27,16 @@ import static lombok.javac.JavacTreeMaker.TreeTag.treeTag;
 import static lombok.javac.JavacTreeMaker.TypeTag.typeTag;
 import static lombok.javac.handlers.JavacHandlerUtil.*;
 
+import java.util.ArrayList;
+
+import com.sun.tools.javac.code.Flags;
+import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.JCTree.JCAnnotation;
 import com.sun.tools.javac.tree.JCTree.JCAssert;
 import com.sun.tools.javac.tree.JCTree.JCAssign;
 import com.sun.tools.javac.tree.JCTree.JCBinary;
 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.JCExpressionStatement;
 import com.sun.tools.javac.tree.JCTree.JCFieldAccess;
@@ -40,13 +45,16 @@ import com.sun.tools.javac.tree.JCTree.JCIf;
 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;
 import com.sun.tools.javac.tree.JCTree.JCParens;
 import com.sun.tools.javac.tree.JCTree.JCStatement;
 import com.sun.tools.javac.tree.JCTree.JCSynchronized;
 import com.sun.tools.javac.tree.JCTree.JCThrow;
 import com.sun.tools.javac.tree.JCTree.JCTry;
+import com.sun.tools.javac.tree.JCTree.JCTypeParameter;
 import com.sun.tools.javac.tree.JCTree.JCVariableDecl;
 import com.sun.tools.javac.util.List;
+import com.sun.tools.javac.util.ListBuffer;
 import com.sun.tools.javac.util.Name;
 
 import lombok.AccessLevel;
@@ -57,83 +65,105 @@ import lombok.core.AnnotationValues;
 import lombok.core.HandlerPriority;
 import lombok.javac.JavacAnnotationHandler;
 import lombok.javac.JavacNode;
-import lombok.javac.handlers.HandleConstructor.SkipIfConstructorExists;
+import lombok.javac.JavacTreeMaker;
 import lombok.spi.Provides;
 
 @Provides
 @HandlerPriority(value = 512) // 2^9; onParameter=@__(@NonNull) has to run first.
 public class HandleNonNull extends JavacAnnotationHandler<NonNull> {
-	private HandleConstructor handleConstructor = new HandleConstructor();
-	
-	@Override public void handle(AnnotationValues<NonNull> annotation, JCAnnotation ast, JavacNode annotationNode) {
-		handleFlagUsage(annotationNode, ConfigurationKeys.NON_NULL_FLAG_USAGE, "@NonNull");
+	private JCMethodDecl createRecordArgslessConstructor(JavacNode typeNode, JavacNode source) {
+		JavacTreeMaker maker = typeNode.getTreeMaker();
 		
-		if (annotationNode.up().getKind() == Kind.FIELD) {
-			// This is meaningless unless the field is used to generate a method (@Setter, @RequiredArgsConstructor, etc),
-			// but in that case those handlers will take care of it. However, we DO check if the annotation is applied to
-			// a primitive, because those handlers trigger on any annotation named @NonNull and we only want the warning
-			// behaviour on _OUR_ 'lombok.NonNull'.
-			
-			try {
-				if (isPrimitive(((JCVariableDecl) annotationNode.up().get()).vartype)) {
-					annotationNode.addWarning("@NonNull is meaningless on a primitive.");
+		java.util.List<JCVariableDecl> fields = new ArrayList<JCVariableDecl>();
+		for (JavacNode child : typeNode.down()) {
+			if (child.getKind() == Kind.FIELD) {
+				JCVariableDecl v = (JCVariableDecl) child.get();
+				if ((v.mods.flags & RECORD) != 0) {
+					fields.add(v);
 				}
-			} catch (Exception ignore) {}
-			
-			return;
+			}
 		}
 		
-		JCMethodDecl declaration;
-		JavacNode paramNode;
+		ListBuffer<JCVariableDecl> params = new ListBuffer<JCVariableDecl>();
 		
-		switch (annotationNode.up().getKind()) {
-		case ARGUMENT:
-			paramNode = annotationNode.up();
-			break;
-		case TYPE_USE:
-			JavacNode typeNode = annotationNode.directUp();
-			paramNode = typeNode.directUp();
-			break;
-		default:
-			return;
+		for (int i = 0; i < fields.size(); i++) {
+			JCVariableDecl arg = fields.get(i);
+			JCModifiers mods = maker.Modifiers(GENERATED_MEMBER | Flags.PARAMETER, arg.mods.annotations);
+			params.append(maker.VarDef(mods, arg.name, arg.vartype, null));
 		}
 		
-		if (paramNode.getKind() != Kind.ARGUMENT) return;
-		try {
-			declaration = (JCMethodDecl) paramNode.up().get();
-		} catch (Exception e) {
-			return;
+		JCModifiers mods = maker.Modifiers(toJavacModifier(AccessLevel.PUBLIC) | COMPACT_RECORD_CONSTRUCTOR, List.<JCAnnotation>nil());
+		JCBlock body = maker.Block(0L, List.<JCStatement>nil());
+		JCMethodDecl constr = maker.MethodDef(mods, typeNode.toName("<init>"), null, List.<JCTypeParameter>nil(), params.toList(), List.<JCExpression>nil(), body, null);
+		return recursiveSetGeneratedBy(constr, source);
+	}
+	
+	/**
+	 * If the provided typeNode is a record, returns the compact constructor (there should only be one, but if the file is
+	 * not semantically sound there might be more). If the only one in existence is the default auto-generated one, it is removed,
+	 * a new explicit one is created, and that one is returned in a list.
+	 * 
+	 * Otherwise, an empty list is returned.
+	 */
+	private List<JCMethodDecl> addCompactConstructorIfNeeded(JavacNode typeNode, JavacNode source) {
+		List<JCMethodDecl> answer = List.nil();
+		
+		if (typeNode == null || !(typeNode.get() instanceof JCClassDecl)) return answer;
+		
+		JCClassDecl cDecl = (JCClassDecl) typeNode.get();
+		if ((cDecl.mods.flags & RECORD) == 0) return answer;
+		
+		ListBuffer<JCTree> newDefs = new ListBuffer<JCTree>();
+		boolean generateConstructor = false;
+		
+		for (JCTree def : cDecl.defs) {
+			boolean remove = false;
+			if (def instanceof JCMethodDecl) {
+				JCMethodDecl md = (JCMethodDecl) def;
+				if (md.name.contentEquals("<init>")) {
+					if ((md.mods.flags & Flags.GENERATEDCONSTR) != 0) {
+						remove = true;
+						generateConstructor = true;
+					} else {
+						if (!isTolerate(typeNode, md)) {
+							if ((md.mods.flags & COMPACT_RECORD_CONSTRUCTOR) != 0) {
+								generateConstructor = false;
+								answer = answer.prepend(md);
+							}
+						}
+					}
+				}
+			}
+			if (!remove) newDefs.append(def);
 		}
 		
-		if (declaration.body == null) {
-			// This used to be a warning, but as @NonNull also has a documentary purpose, better to not warn about this. Since 1.16.7
-			return;
+		if (generateConstructor) {
+			cDecl.defs = newDefs.toList();
+			JCMethodDecl ctr = createRecordArgslessConstructor(typeNode, source);
+			injectMethod(typeNode, ctr);
+			answer = answer.prepend(ctr);
 		}
 		
+		return answer;
+	}
+	
+	private void addNullCheckIfNeeded(JCMethodDecl method, JavacNode paramNode, JavacNode source) {
 		// Possibly, if 'declaration instanceof ConstructorDeclaration', fetch declaration.constructorCall, search it for any references to our parameter,
 		// and if they exist, create a new method in the class: 'private static <T> T lombok$nullCheck(T expr, String msg) {if (expr == null) throw NPE; return expr;}' and
 		// wrap all references to it in the super/this to a call to this method.
 		
-		JCStatement nullCheck = recursiveSetGeneratedBy(generateNullCheck(annotationNode.getTreeMaker(), paramNode, annotationNode), annotationNode);
+		JCStatement nullCheck = recursiveSetGeneratedBy(generateNullCheck(source.getTreeMaker(), paramNode, source), source);
 		
 		if (nullCheck == null) {
 			// @NonNull applied to a primitive. Kinda pointless. Let's generate a warning.
-			annotationNode.addWarning("@NonNull is meaningless on a primitive.");
+			source.addWarning("@NonNull is meaningless on a primitive.");
 			return;
 		}
 		
-		List<JCStatement> statements = declaration.body.stats;
+		List<JCStatement> statements = method.body.stats;
 		
 		String expectedName = paramNode.getName();
 		
-		JavacNode typeNode = upToTypeNode(annotationNode);
-		
-		if ((declaration.mods.flags & RECORD) != 0) {
-			if (!lombokConstructorExists(typeNode)) {
-				handleConstructor.generateAllArgsConstructor(typeNode, AccessLevel.PUBLIC, null, SkipIfConstructorExists.NO, annotationNode);
-			}
-		}
-		
 		/* Abort if the null check is already there, delving into try and synchronized statements */ {
 			List<JCStatement> stats = statements;
 			int idx = 0;
@@ -169,8 +199,69 @@ public class HandleNonNull extends JavacAnnotationHandler<NonNull> {
 		
 		List<JCStatement> newList = tail.prepend(nullCheck);
 		for (JCStatement stat : head) newList = newList.prepend(stat);
-		declaration.body.stats = newList;
-		annotationNode.getAst().setChanged();
+		method.body.stats = newList;
+		source.getAst().setChanged();
+	}
+	
+	@Override public void handle(AnnotationValues<NonNull> annotation, JCAnnotation ast, JavacNode annotationNode) {
+		handleFlagUsage(annotationNode, ConfigurationKeys.NON_NULL_FLAG_USAGE, "@NonNull");
+		if (annotationNode.up().getKind() == Kind.FIELD) {
+			// This is meaningless unless the field is used to generate a method (@Setter, @RequiredArgsConstructor, etc),
+			// but in that case those handlers will take care of it. However, we DO check if the annotation is applied to
+			// a primitive, because those handlers trigger on any annotation named @NonNull and we only want the warning
+			// behaviour on _OUR_ 'lombok.NonNull'.
+			
+			try {
+				if (isPrimitive(((JCVariableDecl) annotationNode.up().get()).vartype)) {
+					annotationNode.addWarning("@NonNull is meaningless on a primitive.");
+				}
+			} catch (Exception ignore) {}
+			
+			JCVariableDecl fDecl = (JCVariableDecl) annotationNode.up().get();
+			if ((fDecl.mods.flags & RECORD) != 0) {
+				List<JCMethodDecl> compactConstructors = addCompactConstructorIfNeeded(annotationNode.up().up(), annotationNode);
+				for (JCMethodDecl ctr : compactConstructors) {
+					addNullCheckIfNeeded(ctr, annotationNode.up(), annotationNode);
+				}
+			}
+			return;
+		}
+		
+		JCMethodDecl declaration;
+		JavacNode paramNode;
+		
+		switch (annotationNode.up().getKind()) {
+		case ARGUMENT:
+			paramNode = annotationNode.up();
+			break;
+		case TYPE_USE:
+			JavacNode typeNode = annotationNode.directUp();
+			paramNode = typeNode.directUp();
+			break;
+		default:
+			return;
+		}
+		
+		if (paramNode.getKind() != Kind.ARGUMENT) return;
+		try {
+			declaration = (JCMethodDecl) paramNode.up().get();
+		} catch (Exception e) {
+			return;
+		}
+		
+		if (declaration.body == null) {
+			// This used to be a warning, but as @NonNull also has a documentary purpose, better to not warn about this. Since 1.16.7
+			return;
+		}
+		
+		if ((declaration.mods.flags & (GENERATED_MEMBER | COMPACT_RECORD_CONSTRUCTOR)) != 0) {
+			// The 'real' annotations are on the `record Foo(@NonNull Obj x)` part and we just see these
+			// syntax-sugared over. We deal with it on the field declaration variant, as those are always there,
+			// not dependent on whether you write out the compact constructor or not.
+			return;
+		}
+		
+		addNullCheckIfNeeded(declaration, paramNode, annotationNode);
 	}
 	
 	public boolean isNullCheck(JCStatement stat) {
diff --git a/src/core/lombok/javac/handlers/JavacHandlerUtil.java b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
index ac00fb1f..32eae4e9 100644
--- a/src/core/lombok/javac/handlers/JavacHandlerUtil.java
+++ b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
@@ -820,30 +820,6 @@ public class JavacHandlerUtil {
 		return MemberExistsResult.NOT_EXISTS;
 	}
 	
-	/**
-	 * Checks if there is at least one constructor that is generated by lombok.
-	 * 
-	 * @param node Any node that represents the Type (TypeDeclaration) to look in, or any child node thereof.
-	 */
-	public static boolean lombokConstructorExists(JavacNode node) {
-		node = upToTypeNode(node);
-		
-		if (node != null && node.get() instanceof JCClassDecl) {
-			for (JCTree def : ((JCClassDecl) node.get()).defs) {
-				if (def instanceof JCMethodDecl) {
-					JCMethodDecl md = (JCMethodDecl) def;
-					if (md.name.contentEquals("<init>")) {
-						if ((md.mods.flags & Flags.GENERATEDCONSTR) != 0) continue;
-						if (isTolerate(node, md)) continue;
-						if (getGeneratedBy(def) != null) return true;
-					}
-				}
-			}
-		}
-		
-		return false;
-	}
-	
 	public static boolean isConstructorCall(final JCStatement statement) {
 		if (!(statement instanceof JCExpressionStatement)) return false;
 		JCExpression expr = ((JCExpressionStatement) statement).expr;