[fixes #3366] Handle @NonNull type annotations on record components

17 files changed, 110 insertions, 81 deletions

diff --git a/src/core/lombok/eclipse/handlers/HandleNonNull.java b/src/core/lombok/eclipse/handlers/HandleNonNull.java
index 365eef33..0138c7fb 100644
--- a/src/core/lombok/eclipse/handlers/HandleNonNull.java
+++ b/src/core/lombok/eclipse/handlers/HandleNonNull.java
@@ -65,7 +65,6 @@ import lombok.core.AST.Kind;
 import lombok.core.AnnotationValues;
 import lombok.core.HandlerPriority;
 import lombok.eclipse.EcjAugments;
-import lombok.eclipse.EclipseAST;
 import lombok.eclipse.EclipseAnnotationHandler;
 import lombok.eclipse.EclipseNode;
 import lombok.spi.Provides;
@@ -221,9 +220,16 @@ public class HandleNonNull extends EclipseAnnotationHandler<NonNull> {
 		// As we need both for @NonNull we reset the handled flag during diet parse.
 		
 		if (!annotationNode.isCompleteParse()) {
-			if (annotationNode.up().getKind() == Kind.FIELD) {
+			final EclipseNode node;
+			if (annotationNode.up().getKind() == Kind.TYPE_USE) {
+				node = annotationNode.directUp().directUp();
+			} else {
+				node = annotationNode.up();
+			}
+			
+			if (node.getKind() == Kind.FIELD) {
 				//Check if this is a record and we need to generate the compact form constructor.
-				EclipseNode typeNode = annotationNode.up().up();
+				EclipseNode typeNode = node.up();
 				if (typeNode.getKind() == Kind.TYPE) {
 					if (isRecord(typeNode)) {
 						addCompactConstructorIfNeeded(typeNode, annotationNode);
@@ -251,16 +257,23 @@ public class HandleNonNull extends EclipseAnnotationHandler<NonNull> {
 	private void handle0(Annotation ast, EclipseNode annotationNode, boolean force) {
 		handleFlagUsage(annotationNode, ConfigurationKeys.NON_NULL_FLAG_USAGE, "@NonNull");
 		
-		if (annotationNode.up().getKind() == Kind.FIELD) {
+		final EclipseNode node;
+		if (annotationNode.up().getKind() == Kind.TYPE_USE) {
+			node = annotationNode.directUp().directUp();
+		} else {
+			node = annotationNode.up();
+		}
+		
+		if (node.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'.
-			EclipseNode fieldNode = annotationNode.up();
+			EclipseNode fieldNode = node;
 			EclipseNode typeNode = fieldNode.up();
 			
 			try {
-				if (isPrimitive(((AbstractVariableDeclaration) annotationNode.up().get()).type)) {
+				if (isPrimitive(((AbstractVariableDeclaration) node.get()).type)) {
 					annotationNode.addWarning("@NonNull is meaningless on a primitive.");
 					return;
 				}
@@ -279,33 +292,14 @@ public class HandleNonNull extends EclipseAnnotationHandler<NonNull> {
 			return;
 		}
 		
+		if (node.getKind() != Kind.ARGUMENT) return;
+		
 		Argument param;
-		EclipseNode paramNode;
 		AbstractMethodDeclaration declaration;
 		
-		switch (annotationNode.up().getKind()) {
-		case ARGUMENT:
-			paramNode = annotationNode.up();
-			break;
-		case TYPE_USE:
-			EclipseNode typeNode = annotationNode.directUp();
-			boolean ok = false;
-			ASTNode astNode = typeNode.get();
-			if (astNode instanceof TypeReference) {
-				Annotation[] anns = EclipseAST.getTopLevelTypeReferenceAnnotations((TypeReference) astNode);
-				if (anns == null) return;
-				for (Annotation ann : anns) if (ast == ann) ok = true;
-			}
-			if (!ok) return;
-			paramNode = typeNode.directUp();
-			break;
-		default:
-			return;
-		}
-		
 		try {
-			param = (Argument) paramNode.get();
-			declaration = (AbstractMethodDeclaration) paramNode.up().get();
+			param = (Argument) node.get();
+			declaration = (AbstractMethodDeclaration) node.up().get();
 		} catch (Exception e) {
 			return;
 		}
@@ -318,7 +312,7 @@ public class HandleNonNull extends EclipseAnnotationHandler<NonNull> {
 		}
 		
 		addNullCheckIfNeeded(declaration, param, annotationNode);
-		paramNode.up().rebuild();
+		node.up().rebuild();
 	}
 	
 	private void addNullCheckIfNeeded(AbstractMethodDeclaration declaration, AbstractVariableDeclaration param, EclipseNode annotationNode) {
diff --git a/src/core/lombok/javac/handlers/HandleNonNull.java b/src/core/lombok/javac/handlers/HandleNonNull.java
index 271bedbb..e817916d 100644
--- a/src/core/lombok/javac/handlers/HandleNonNull.java
+++ b/src/core/lombok/javac/handlers/HandleNonNull.java
@@ -218,48 +218,43 @@ public class HandleNonNull extends JavacAnnotationHandler<NonNull> {
 	
 	@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) {
+		
+		final JavacNode node;
+		if (annotationNode.up().getKind() == Kind.TYPE_USE) {
+			node = annotationNode.directUp().directUp();
+		} else {
+			node = annotationNode.up();
+		}
+		
+		if (node.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)) {
+				if (isPrimitive(((JCVariableDecl) node.get()).vartype)) {
 					annotationNode.addWarning("@NonNull is meaningless on a primitive.");
 				}
 			} catch (Exception ignore) {}
 			
-			JCVariableDecl fDecl = (JCVariableDecl) annotationNode.up().get();
+			JCVariableDecl fDecl = (JCVariableDecl) node.get();
 			if ((fDecl.mods.flags & RECORD) != 0) {
 				// well, these kinda double as parameters (of the compact constructor), so we do some work here.
 				
-				List<JCMethodDecl> compactConstructors = addCompactConstructorIfNeeded(annotationNode.up().up(), annotationNode);
+				List<JCMethodDecl> compactConstructors = addCompactConstructorIfNeeded(node.up(), annotationNode);
 				for (JCMethodDecl ctr : compactConstructors) {
-					addNullCheckIfNeeded(ctr, annotationNode.up(), annotationNode);
+					addNullCheckIfNeeded(ctr, node, annotationNode);
 				}
 			}
 			return;
 		}
 		
-		JCMethodDecl declaration;
-		JavacNode paramNode;
+		if (node.getKind() != Kind.ARGUMENT) return;
 		
-		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;
+		JCMethodDecl declaration;
 		try {
-			declaration = (JCMethodDecl) paramNode.up().get();
+			declaration = (JCMethodDecl) node.up().get();
 		} catch (Exception e) {
 			return;
 		}
@@ -276,7 +271,7 @@ public class HandleNonNull extends JavacAnnotationHandler<NonNull> {
 			return;
 		}
 		
-		addNullCheckIfNeeded(declaration, paramNode, annotationNode);
+		addNullCheckIfNeeded(declaration, node, annotationNode);
 	}
 	
 	public boolean isNullCheck(JCStatement stat) {
diff --git a/test/transform/resource/after-delombok/NonNullOnRecord2.java b/test/transform/resource/after-delombok/NonNullOnRecordExistingConstructor.java
index d3021350..e3fdc2fc 100644
--- a/test/transform/resource/after-delombok/NonNullOnRecord2.java
+++ b/test/transform/resource/after-delombok/NonNullOnRecordExistingConstructor.java
@@ -1,7 +1,7 @@
 // version 16:
 import lombok.NonNull;
-record NonNullOnRecord2(@NonNull String a) {
-	public NonNullOnRecord2 {
+public record NonNullOnRecordExistingConstructor(@NonNull String a) {
+	public NonNullOnRecordExistingConstructor {
 		if (a == null) {
 			throw new java.lang.NullPointerException("a is marked non-null but is null");
 		}
diff --git a/test/transform/resource/after-delombok/NonNullOnRecord3.java b/test/transform/resource/after-delombok/NonNullOnRecordExistingSetter.java
index 62b385bc..11aaca49 100644
--- a/test/transform/resource/after-delombok/NonNullOnRecord3.java
+++ b/test/transform/resource/after-delombok/NonNullOnRecordExistingSetter.java
@@ -1,7 +1,7 @@
 // version 14:
 import lombok.NonNull;
-public record NonNullOnRecord3(@NonNull String a) {
-	public NonNullOnRecord3(String a) {
+public record NonNullOnRecordExistingSetter(@NonNull String a) {
+	public NonNullOnRecordExistingSetter(String a) {
 		this.a = a;
 	}
 	public void method(@NonNull String param) {
diff --git a/test/transform/resource/after-delombok/NonNullOnRecord.java b/test/transform/resource/after-delombok/NonNullOnRecordSimple.java
index 465c30db..2a772c93 100644
--- a/test/transform/resource/after-delombok/NonNullOnRecord.java
+++ b/test/transform/resource/after-delombok/NonNullOnRecordSimple.java
@@ -1,8 +1,8 @@
 // version 16:
 import lombok.NonNull;
-public record NonNullOnRecord(@NonNull String a, @NonNull String b) {
+public record NonNullOnRecordSimple(@NonNull String a, @NonNull String b) {
 	@java.lang.SuppressWarnings("all")
-	public NonNullOnRecord {
+	public NonNullOnRecordSimple {
 		if (a == null) {
 			throw new java.lang.NullPointerException("a is marked non-null but is null");
 		}
diff --git a/test/transform/resource/after-delombok/NonNullOnRecordTypeUse.java b/test/transform/resource/after-delombok/NonNullOnRecordTypeUse.java
new file mode 100644
index 00000000..e5b78f1c
--- /dev/null
+++ b/test/transform/resource/after-delombok/NonNullOnRecordTypeUse.java
@@ -0,0 +1,13 @@
+// version 16:
+import lombok.NonNull;
+public record NonNullOnRecordTypeUse(@NonNull int[] a, int @NonNull [] b, int[] @NonNull [] c) {
+	@java.lang.SuppressWarnings("all")
+	public NonNullOnRecordTypeUse {
+		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");
+		}
+	}
+}
diff --git a/test/transform/resource/after-ecj/NonNullOnRecord2.java b/test/transform/resource/after-ecj/NonNullOnRecordExistingConstructor.java
index 5820d453..4323113f 100644
--- a/test/transform/resource/after-ecj/NonNullOnRecord2.java
+++ b/test/transform/resource/after-ecj/NonNullOnRecordExistingConstructor.java
@@ -1,8 +1,8 @@
 // version 14:
 import lombok.NonNull;
-record NonNullOnRecord2(String a) {
+public record NonNullOnRecordExistingConstructor(String a) {
 /* Implicit */  private final String a;
-  public NonNullOnRecord2(@NonNull String a) {
+  public NonNullOnRecordExistingConstructor(@NonNull String a) {
     super();
     if ((a == null))
         {
diff --git a/test/transform/resource/after-ecj/NonNullOnRecord3.java b/test/transform/resource/after-ecj/NonNullOnRecordExistingSetter.java
index 37f0afcf..1bdaf29f 100644
--- a/test/transform/resource/after-ecj/NonNullOnRecord3.java
+++ b/test/transform/resource/after-ecj/NonNullOnRecordExistingSetter.java
@@ -1,8 +1,8 @@
 // version 19:
 import lombok.NonNull;
-public record NonNullOnRecord3(String a) {
+public record NonNullOnRecordExistingSetter(String a) {
 /* Implicit */  private final String a;
-  public NonNullOnRecord3(String a) {
+  public NonNullOnRecordExistingSetter(String a) {
     super();
     this.a = a;
   }
diff --git a/test/transform/resource/after-ecj/NonNullOnRecord.java b/test/transform/resource/after-ecj/NonNullOnRecordSimple.java
index d80e243b..ee1510ea 100644
--- a/test/transform/resource/after-ecj/NonNullOnRecord.java
+++ b/test/transform/resource/after-ecj/NonNullOnRecordSimple.java
@@ -1,9 +1,9 @@
 // version 14:
 import lombok.NonNull;
-public record NonNullOnRecord(String a, String b) {
+public record NonNullOnRecordSimple(String a, String b) {
 /* Implicit */  private final String a;
 /* Implicit */  private final String b;
-  public @java.lang.SuppressWarnings("all") NonNullOnRecord(@NonNull String a, @NonNull String b) {
+  public @java.lang.SuppressWarnings("all") NonNullOnRecordSimple(@NonNull String a, @NonNull String b) {
     super();
     if ((a == null))
         {
diff --git a/test/transform/resource/after-ecj/NonNullOnRecordTypeUse.java b/test/transform/resource/after-ecj/NonNullOnRecordTypeUse.java
new file mode 100644
index 00000000..4f040c2e
--- /dev/null
+++ b/test/transform/resource/after-ecj/NonNullOnRecordTypeUse.java
@@ -0,0 +1,21 @@
+// version 14:
+import lombok.NonNull;
+public record NonNullOnRecordTypeUse(int a, int b, int c) {
+/* Implicit */  private final int[] a;
+/* Implicit */  private final int @NonNull [] b;
+/* Implicit */  private final int[] @NonNull [] c;
+  public @java.lang.SuppressWarnings("all") NonNullOnRecordTypeUse(@NonNull int[] a, int @NonNull [] b, int[] @NonNull [] c) {
+    super();
+    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;
+    this.c = c;
+  }
+}
diff --git a/test/transform/resource/before/NonNullOnRecord.java b/test/transform/resource/before/NonNullOnRecord.java
deleted file mode 100644
index ba6121a6..00000000
--- a/test/transform/resource/before/NonNullOnRecord.java
+++ /dev/null
@@ -1,6 +0,0 @@
-// version 14:
-
-import lombok.NonNull;
-
-public record NonNullOnRecord(@NonNull String a, @NonNull String b) {
-}
\ No newline at end of file
diff --git a/test/transform/resource/before/NonNullOnRecord2.java b/test/transform/resource/before/NonNullOnRecord2.java
deleted file mode 100644
index 3a4eacd4..00000000
--- a/test/transform/resource/before/NonNullOnRecord2.java
+++ /dev/null
@@ -1,9 +0,0 @@
-// 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/NonNullOnRecordExistingConstructor.java b/test/transform/resource/before/NonNullOnRecordExistingConstructor.java
new file mode 100644
index 00000000..7bab2a8a
--- /dev/null
+++ b/test/transform/resource/before/NonNullOnRecordExistingConstructor.java
@@ -0,0 +1,9 @@
+// version 14:
+
+import lombok.NonNull;
+
+public record NonNullOnRecordExistingConstructor(@NonNull String a) {
+	public NonNullOnRecordExistingConstructor {
+		System.out.println("Hello");
+	}
+}
\ No newline at end of file
diff --git a/test/transform/resource/before/NonNullOnRecord3.java b/test/transform/resource/before/NonNullOnRecordExistingSetter.java
index 24198002..dd7b15d8 100644
--- a/test/transform/resource/before/NonNullOnRecord3.java
+++ b/test/transform/resource/before/NonNullOnRecordExistingSetter.java
@@ -2,8 +2,8 @@
 
 import lombok.NonNull;
 
-public record NonNullOnRecord3(@NonNull String a) {
-	public NonNullOnRecord3(String a) {
+public record NonNullOnRecordExistingSetter(@NonNull String a) {
+	public NonNullOnRecordExistingSetter(String a) {
 		this.a = a;
 	}
 	
diff --git a/test/transform/resource/before/NonNullOnRecordSimple.java b/test/transform/resource/before/NonNullOnRecordSimple.java
new file mode 100644
index 00000000..0d0a1f9f
--- /dev/null
+++ b/test/transform/resource/before/NonNullOnRecordSimple.java
@@ -0,0 +1,6 @@
+// version 14:
+
+import lombok.NonNull;
+
+public record NonNullOnRecordSimple(@NonNull String a, @NonNull String b) {
+}
\ No newline at end of file
diff --git a/test/transform/resource/before/NonNullOnRecordTypeUse.java b/test/transform/resource/before/NonNullOnRecordTypeUse.java
new file mode 100644
index 00000000..5f89695d
--- /dev/null
+++ b/test/transform/resource/before/NonNullOnRecordTypeUse.java
@@ -0,0 +1,6 @@
+// version 14:
+
+import lombok.NonNull;
+
+public record NonNullOnRecordTypeUse(@NonNull int [] a, int @NonNull [] b, int [] @NonNull [] c) {
+}
\ No newline at end of file
diff --git a/test/transform/resource/before/RecordNonNull b/test/transform/resource/before/RecordNonNull
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/test/transform/resource/before/RecordNonNull