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

2 files changed, 41 insertions, 52 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) {