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

1 files changed, 17 insertions, 22 deletions

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) {