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

11 files changed, 200 insertions, 95 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;
diff --git a/test/core/src/lombok/RunTestsViaDelombok.java b/test/core/src/lombok/RunTestsViaDelombok.java
index 13acdf55..46eb54a0 100644
--- a/test/core/src/lombok/RunTestsViaDelombok.java
+++ b/test/core/src/lombok/RunTestsViaDelombok.java
@@ -120,6 +120,7 @@ public class RunTestsViaDelombok extends AbstractRunTests {
 				@Override public void scan(JCTree tree) {
 					if (tree == null) return;
 					if (tree instanceof JCMethodDecl && (((JCMethodDecl) tree).mods.flags & Flags.GENERATEDCONSTR) != 0) return;
+					astContext.push(tree);
 					try {
 						if (tree instanceof JCModifiers) return;
 						
@@ -139,14 +140,19 @@ public class RunTestsViaDelombok extends AbstractRunTests {
 							if ("super".equals("" + ((JCIdent) tree).name)) check = false;
 						}
 						
+						if (tree instanceof JCVariableDecl && (((JCVariableDecl) tree).mods.flags & Javac.GENERATED_MEMBER) != 0) return;
+						
 						if (check && tree.pos == -1) fail(craftFailMsg("Start", astContext));
+						
 						if (check && Javac.getEndPosition(tree, unit) == -1) {
 							fail(craftFailMsg("End", astContext));
 						}
 					} finally {
-						astContext.push(tree);
-						super.scan(tree);
-						astContext.pop();
+						try {
+							super.scan(tree);
+						} finally {
+							astContext.pop();
+						}
 					}
 				}
 				
diff --git a/test/transform/resource/after-delombok/NonNullExistingConstructorOnRecord.java b/test/transform/resource/after-delombok/NonNullExistingConstructorOnRecord.java
index 02b66dd1..15e0aa66 100644
--- a/test/transform/resource/after-delombok/NonNullExistingConstructorOnRecord.java
+++ b/test/transform/resource/after-delombok/NonNullExistingConstructorOnRecord.java
@@ -8,14 +8,12 @@ record NonNullOnRecord(@NonNull String a, @NonNull String b) {
 		}
 	}
 	@java.lang.SuppressWarnings("all")
-	public NonNullOnRecord(@NonNull final String a, @NonNull final String b) {
+	public NonNullOnRecord {
 		if (a == null) {
 			throw new java.lang.NullPointerException("a is marked non-null but is null");
 		}
 		if (b == null) {
 			throw new java.lang.NullPointerException("b is marked non-null but is null");
 		}
-		this.a = a;
-		this.b = b;
 	}
 }
diff --git a/test/transform/resource/after-delombok/NonNullOnRecord.java b/test/transform/resource/after-delombok/NonNullOnRecord.java
index 2ff4a8f7..7acfab36 100644
--- a/test/transform/resource/after-delombok/NonNullOnRecord.java
+++ b/test/transform/resource/after-delombok/NonNullOnRecord.java
@@ -1,21 +1,13 @@
 // version 14:
 import lombok.NonNull;
 record NonNullOnRecord(@NonNull String a, @NonNull String b) {
-	public void method(@NonNull String param) {
-		if (param == null) {
-			throw new java.lang.NullPointerException("param is marked non-null but is null");
-		}
-		String asd = "a";
-	}
 	@java.lang.SuppressWarnings("all")
-	public NonNullOnRecord(@NonNull final String a, @NonNull final String b) {
+	public NonNullOnRecord {
 		if (a == null) {
 			throw new java.lang.NullPointerException("a is marked non-null but is null");
 		}
 		if (b == null) {
 			throw new java.lang.NullPointerException("b is marked non-null but is null");
 		}
-		this.a = a;
-		this.b = b;
 	}
 }
diff --git a/test/transform/resource/after-delombok/NonNullOnRecord2.java b/test/transform/resource/after-delombok/NonNullOnRecord2.java
new file mode 100644
index 00000000..a166190c
--- /dev/null
+++ b/test/transform/resource/after-delombok/NonNullOnRecord2.java
@@ -0,0 +1,10 @@
+// version 14:
+import lombok.NonNull;
+record NonNullOnRecord2(@NonNull String a) {
+	public NonNullOnRecord2 {
+		if (a == null) {
+			throw new java.lang.NullPointerException("a is marked non-null but is null");
+		}
+		System.out.println("Hello");
+	}
+}
diff --git a/test/transform/resource/after-delombok/NonNullOnRecord3.java b/test/transform/resource/after-delombok/NonNullOnRecord3.java
new file mode 100644
index 00000000..b1ecf6c0
--- /dev/null
+++ b/test/transform/resource/after-delombok/NonNullOnRecord3.java
@@ -0,0 +1,13 @@
+// version 14:
+import lombok.NonNull;
+record NonNullOnRecord3(@NonNull String a) {
+	public NonNullOnRecord3(String a) {
+		this.a = a;
+	}
+	public void method(@NonNull String param) {
+		if (param == null) {
+			throw new java.lang.NullPointerException("param is marked non-null but is null");
+		}
+		String asd = "a";
+	}
+}
diff --git a/test/transform/resource/before/NonNullOnRecord.java b/test/transform/resource/before/NonNullOnRecord.java
index 223f8cbb..56a96433 100644
--- a/test/transform/resource/before/NonNullOnRecord.java
+++ b/test/transform/resource/before/NonNullOnRecord.java
@@ -3,7 +3,4 @@
 import lombok.NonNull;
 
 record NonNullOnRecord(@NonNull String a, @NonNull String b) {
-	public void method(@NonNull String param) {
-		String asd = "a";
-	}
 }
\ No newline at end of file
diff --git a/test/transform/resource/before/NonNullOnRecord2.java b/test/transform/resource/before/NonNullOnRecord2.java
new file mode 100644
index 00000000..3a4eacd4
--- /dev/null
+++ b/test/transform/resource/before/NonNullOnRecord2.java
@@ -0,0 +1,9 @@
+// version 14:
+
+import lombok.NonNull;
+
+record NonNullOnRecord2(@NonNull String a) {
+	public NonNullOnRecord2 {
+		System.out.println("Hello");
+	}
+}
\ No newline at end of file
diff --git a/test/transform/resource/before/NonNullOnRecord3.java b/test/transform/resource/before/NonNullOnRecord3.java
new file mode 100644
index 00000000..88870192
--- /dev/null
+++ b/test/transform/resource/before/NonNullOnRecord3.java
@@ -0,0 +1,13 @@
+// version 14:
+
+import lombok.NonNull;
+
+record NonNullOnRecord3(@NonNull String a) {
+	public NonNullOnRecord3(String a) {
+		this.a = a;
+	}
+	
+	public void method(@NonNull String param) {
+		String asd = "a";
+	}
+}
\ No newline at end of file
diff --git a/website/templates/features/NonNull.html b/website/templates/features/NonNull.html
index 67cc262c..57aa62fe 100644
--- a/website/templates/features/NonNull.html
+++ b/website/templates/features/NonNull.html
@@ -7,11 +7,11 @@
 	
 	<@f.overview>
 		<p>
-			You can use <code>@NonNull</code> on the parameter of a method or constructor to have lombok generate a null-check statement for you.
+			You can use <code>@NonNull</code> on a record component, or a parameter of a method or constructor. This will cause to lombok generate a null-check statement for you.
 		</p><p>
-			Lombok has always treated various annotations generally named <code>@NonNull</code> on a field as a signal to generate a null-check if lombok generates an entire method or constructor for you, via for example <a href="/features/Data"><code>@Data</code></a>. Now, however, using lombok's own <code>@lombok.NonNull</code> on a parameter results in the insertion of just the null-check statement inside your own method or constructor.
+			Lombok has always treated various annotations generally named <code>@NonNull</code> on a field as a signal to generate a null-check if lombok generates an entire method or constructor for you, via for example <a href="/features/Data"><code>@Data</code></a>. However, using lombok's own <code>@lombok.NonNull</code> on a parameter or record component results in the insertion of the null-check at the top of that method.
 		</p><p>
-			The null-check looks like <code>if (param == null) throw new NullPointerException("param is marked @NonNull but is null");</code> and will be inserted at the very top of your method. For constructors, the null-check will be inserted immediately following any explicit <code>this()</code> or <code>super()</code> calls.
+			The null-check looks like <code>if (param == null) throw new NullPointerException("param is marked @NonNull but is null");</code> and will be inserted at the very top of your method. For constructors, the null-check will be inserted immediately following any explicit <code>this()</code> or <code>super()</code> calls. For record components, the null-check will be inserted in the 'compact constructor' (the one that has no argument list at all), which will be generated if you have no constructor. If you have written out the record constructor in long form (with parameters matching your components exactly), then nothing happens - you'd have to annotate the parameters of this long-form constructor instead.
 		</p><p>
 			If a null-check is already present at the top, no additional null-check will be generated.
 		</p>