[fixes #2019] Lombok now properly deals with `@NonNull` specifically on the ‘type use’ of a parameter (and, in case of arrays, on the outermost dimension which is actually the first one listed. Weird corner case of the JLS).

7 files changed, 132 insertions, 20 deletions

diff --git a/src/core/lombok/javac/JavacAST.java b/src/core/lombok/javac/JavacAST.java
index f2901038..f6cd5571 100644
--- a/src/core/lombok/javac/JavacAST.java
+++ b/src/core/lombok/javac/JavacAST.java
@@ -47,15 +47,19 @@ import com.sun.tools.javac.model.JavacElements;
 import com.sun.tools.javac.model.JavacTypes;
 import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.JCTree.JCAnnotation;
+import com.sun.tools.javac.tree.JCTree.JCArrayTypeTree;
 import com.sun.tools.javac.tree.JCTree.JCBlock;
 import com.sun.tools.javac.tree.JCTree.JCCatch;
 import com.sun.tools.javac.tree.JCTree.JCClassDecl;
 import com.sun.tools.javac.tree.JCTree.JCCompilationUnit;
 import com.sun.tools.javac.tree.JCTree.JCExpression;
+import com.sun.tools.javac.tree.JCTree.JCFieldAccess;
+import com.sun.tools.javac.tree.JCTree.JCIdent;
 import com.sun.tools.javac.tree.JCTree.JCMethodDecl;
 import com.sun.tools.javac.tree.JCTree.JCStatement;
 import com.sun.tools.javac.tree.JCTree.JCTry;
 import com.sun.tools.javac.tree.JCTree.JCVariableDecl;
+import com.sun.tools.javac.tree.JCTree.JCWildcard;
 import com.sun.tools.javac.tree.TreeMaker;
 import com.sun.tools.javac.util.Context;
 import com.sun.tools.javac.util.JCDiagnostic.DiagnosticPosition;
@@ -192,6 +196,8 @@ public class JavacAST extends AST<JavacAST, JavacNode, JCTree> {
 			return buildStatementOrExpression(node);
 		case ANNOTATION:
 			return buildAnnotation((JCAnnotation) node, false);
+		case TYPE_USE:
+			return buildTypeUse(node);
 		default:
 			throw new AssertionError("Did not expect: " + kind);
 		}
@@ -233,6 +239,7 @@ public class JavacAST extends AST<JavacAST, JavacNode, JCTree> {
 		if (setAndGetAsHandled(field)) return null;
 		List<JavacNode> childNodes = new ArrayList<JavacNode>();
 		for (JCAnnotation annotation : field.mods.annotations) addIfNotNull(childNodes, buildAnnotation(annotation, true));
+		addIfNotNull(childNodes, buildTypeUse(field.vartype));
 		addIfNotNull(childNodes, buildExpression(field.init));
 		return putInMap(new JavacNode(this, field, childNodes, Kind.FIELD));
 	}
@@ -241,23 +248,62 @@ public class JavacAST extends AST<JavacAST, JavacNode, JCTree> {
 		if (setAndGetAsHandled(local)) return null;
 		List<JavacNode> childNodes = new ArrayList<JavacNode>();
 		for (JCAnnotation annotation : local.mods.annotations) addIfNotNull(childNodes, buildAnnotation(annotation, true));
+		addIfNotNull(childNodes, buildTypeUse(local.vartype));
 		addIfNotNull(childNodes, buildExpression(local.init));
 		return putInMap(new JavacNode(this, local, childNodes, kind));
 	}
 	
-	private static boolean JCTRY_RESOURCES_FIELD_INITIALIZED;
+	private JavacNode buildTypeUse(JCTree typeUse) {
+		if (setAndGetAsHandled(typeUse)) return null;
+		
+		if (typeUse == null) return null;
+		
+		if (typeUse.getClass().getSimpleName().equals("JCAnnotatedType")) {
+			initJcAnnotatedType(typeUse.getClass());
+			Collection<?> anns = Permit.permissiveReadField(Collection.class, JCANNOTATEDTYPE_ANNOTATIONS, typeUse);
+			JCExpression underlying = Permit.permissiveReadField(JCExpression.class, JCANNOTATEDTYPE_UNDERLYINGTYPE, typeUse);
+			
+			List<JavacNode> childNodes = new ArrayList<JavacNode>();
+			if (anns != null) for (Object annotation : anns) if (annotation instanceof JCAnnotation) addIfNotNull(childNodes, buildAnnotation((JCAnnotation) annotation, true));
+			addIfNotNull(childNodes, buildTypeUse(underlying));
+			return putInMap(new JavacNode(this, typeUse, childNodes, Kind.TYPE_USE));
+		}
+		
+		if (typeUse instanceof JCWildcard) {
+			JCTree inner = ((JCWildcard) typeUse).inner;
+			List<JavacNode> childNodes = inner == null ? Collections.<JavacNode>emptyList() : new ArrayList<JavacNode>();
+			if (inner != null) addIfNotNull(childNodes, buildTypeUse(inner));
+			return putInMap(new JavacNode(this, typeUse, childNodes, Kind.TYPE_USE));
+		}
+		
+		if (typeUse instanceof JCArrayTypeTree) {
+			JCTree inner = ((JCArrayTypeTree) typeUse).elemtype;
+			List<JavacNode> childNodes = inner == null ? Collections.<JavacNode>emptyList() : new ArrayList<JavacNode>();
+			if (inner != null) addIfNotNull(childNodes, buildTypeUse(inner));
+			return putInMap(new JavacNode(this, typeUse, childNodes, Kind.TYPE_USE));
+		}
+		
+		if (typeUse instanceof JCFieldAccess) {
+			JCTree inner = ((JCFieldAccess) typeUse).selected;
+			List<JavacNode> childNodes = inner == null ? Collections.<JavacNode>emptyList() : new ArrayList<JavacNode>();
+			if (inner != null) addIfNotNull(childNodes, buildTypeUse(inner));
+			return putInMap(new JavacNode(this, typeUse, childNodes, Kind.TYPE_USE));
+		}
+		
+		if (typeUse instanceof JCIdent) {
+			return putInMap(new JavacNode(this, typeUse, Collections.<JavacNode>emptyList(), Kind.TYPE_USE));
+		}
+		
+		return null;
+	}
+	
+	private static boolean JCTRY_RESOURCES_FIELD_INITIALIZED = false;
 	private static Field JCTRY_RESOURCES_FIELD;
 	
 	@SuppressWarnings("unchecked")
 	private static List<JCTree> getResourcesForTryNode(JCTry tryNode) {
 		if (!JCTRY_RESOURCES_FIELD_INITIALIZED) {
-			try {
-				JCTRY_RESOURCES_FIELD = Permit.getField(JCTry.class, "resources");
-			} catch (NoSuchFieldException ignore) {
-				// Java 1.6 or lower won't have this at all.
-			} catch (Exception ignore) {
-				// Shouldn't happen. Best thing we can do is just carry on and break on try/catch.
-			}
+			JCTRY_RESOURCES_FIELD = Permit.permissiveGetField(JCTry.class, "resources");
 			JCTRY_RESOURCES_FIELD_INITIALIZED = true;
 		}
 		
@@ -271,6 +317,15 @@ public class JavacAST extends AST<JavacAST, JavacNode, JCTree> {
 		return Collections.emptyList();
 	}
 	
+	private static boolean JCANNOTATEDTYPE_FIELDS_INITIALIZED = false;
+	private static Field JCANNOTATEDTYPE_ANNOTATIONS, JCANNOTATEDTYPE_UNDERLYINGTYPE;
+	private static void initJcAnnotatedType(Class<?> context) {
+		if (JCANNOTATEDTYPE_FIELDS_INITIALIZED) return;
+		JCANNOTATEDTYPE_ANNOTATIONS = Permit.permissiveGetField(context, "annotations");
+		JCANNOTATEDTYPE_UNDERLYINGTYPE = Permit.permissiveGetField(context, "underlyingType");
+		JCANNOTATEDTYPE_FIELDS_INITIALIZED = true;
+	}
+	
 	private JavacNode buildTry(JCTry tryNode) {
 		if (setAndGetAsHandled(tryNode)) return null;
 		List<JavacNode> childNodes = new ArrayList<JavacNode>();
@@ -323,8 +378,8 @@ public class JavacAST extends AST<JavacAST, JavacNode, JCTree> {
 	private JavacNode buildStatementOrExpression(JCTree statement) {
 		if (statement == null) return null;
 		if (statement instanceof JCAnnotation) return null;
-		if (statement instanceof JCClassDecl) return buildType((JCClassDecl)statement);
-		if (statement instanceof JCVariableDecl) return buildLocalVar((JCVariableDecl)statement, Kind.LOCAL);
+		if (statement instanceof JCClassDecl) return buildType((JCClassDecl) statement);
+		if (statement instanceof JCVariableDecl) return buildLocalVar((JCVariableDecl) statement, Kind.LOCAL);
 		if (statement instanceof JCTry) return buildTry((JCTry) statement);
 		if (statement.getClass().getSimpleName().equals("JCLambda")) return buildLambda(statement);
 		if (setAndGetAsHandled(statement)) return null;
diff --git a/src/core/lombok/javac/JavacASTAdapter.java b/src/core/lombok/javac/JavacASTAdapter.java
index 6af53e3d..4c1912d8 100644
--- a/src/core/lombok/javac/JavacASTAdapter.java
+++ b/src/core/lombok/javac/JavacASTAdapter.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2009-2012 The Project Lombok Authors.
+ * Copyright (C) 2009-2019 The Project Lombok Authors.
  * 
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -96,6 +96,15 @@ public class JavacASTAdapter implements JavacASTVisitor {
 	@Override public void endVisitLocal(JavacNode localNode, JCVariableDecl local) {}
 	
 	/** {@inheritDoc} */
+	@Override public void visitTypeUse(JavacNode typeUseNode, JCTree typeUse) {}
+	
+	/** {@inheritDoc} */
+	@Override public void visitAnnotationOnTypeUse(JCTree typeUse, JavacNode annotationNode, JCAnnotation annotation) {}
+	
+	/** {@inheritDoc} */
+	@Override public void endVisitTypeUse(JavacNode typeUseNode, JCTree typeUse) {}
+	
+	/** {@inheritDoc} */
 	@Override public void visitStatement(JavacNode statementNode, JCTree statement) {}
 	
 	/** {@inheritDoc} */
diff --git a/src/core/lombok/javac/JavacASTVisitor.java b/src/core/lombok/javac/JavacASTVisitor.java
index d4f8f731..9b67dda3 100644
--- a/src/core/lombok/javac/JavacASTVisitor.java
+++ b/src/core/lombok/javac/JavacASTVisitor.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2009-2012 The Project Lombok Authors.
+ * Copyright (C) 2009-2019 The Project Lombok Authors.
  * 
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -89,6 +89,13 @@ public interface JavacASTVisitor {
 	void endVisitLocal(JavacNode localNode, JCVariableDecl local);
 	
 	/**
+	 * Visits a node that represents a type reference. Anything from {@code int} to {@code T} to {@code foo.pkg.Bar<T>.Baz<?> @Ann []}.
+	 */
+	void visitTypeUse(JavacNode typeUseNode, JCTree typeUse);
+	void visitAnnotationOnTypeUse(JCTree typeUse, JavacNode annotationNode, JCAnnotation annotation);
+	void endVisitTypeUse(JavacNode typeUseNode, JCTree typeUse);
+	
+	/**
 	 * Visits a statement that isn't any of the other visit methods (e.g. JCClassDecl).
 	 * The statement object is guaranteed to be either a JCStatement or a JCExpression.
 	 */
@@ -261,6 +268,21 @@ public interface JavacASTVisitor {
 			print("</LOCAL %s %s>", local.vartype, local.name);
 		}
 		
+		@Override public void visitTypeUse(JavacNode node, JCTree typeUse) {
+			print("<TYPE %s>", typeUse.getClass());
+			indent++;
+			print("%s", typeUse);
+		}
+		
+		@Override public void visitAnnotationOnTypeUse(JCTree typeUse, JavacNode node, JCAnnotation annotation) {
+			print("<ANNOTATION: %s />", annotation);
+		}
+		
+		@Override public void endVisitTypeUse(JavacNode node, JCTree typeUse) {
+			indent--;
+			print("</TYPE %s>", typeUse.getClass());
+		}
+		
 		@Override public void visitStatement(JavacNode node, JCTree statement) {
 			print("<%s>", statement.getClass());
 			indent++;
diff --git a/src/core/lombok/javac/JavacNode.java b/src/core/lombok/javac/JavacNode.java
index f119f1f9..191ab3c0 100644
--- a/src/core/lombok/javac/JavacNode.java
+++ b/src/core/lombok/javac/JavacNode.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2009-2018 The Project Lombok Authors.
+ * Copyright (C) 2009-2019 The Project Lombok Authors.
  * 
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -145,10 +145,18 @@ public class JavacNode extends lombok.core.LombokNode<JavacAST, JavacNode, JCTre
 			case LOCAL:
 				visitor.visitAnnotationOnLocal((JCVariableDecl) up().get(), this, (JCAnnotation) get());
 				break;
+			case TYPE_USE:
+				visitor.visitAnnotationOnTypeUse(up().get(), this, (JCAnnotation) get());
+				break;
 			default:
 				throw new AssertionError("Annotion not expected as child of a " + up().getKind());
 			}
 			break;
+		case TYPE_USE:
+			visitor.visitTypeUse(this, get());
+			ast.traverseChildren(visitor, this);
+			visitor.endVisitTypeUse(this, get());
+			break;
 		default:
 			throw new AssertionError("Unexpected kind during node traversal: " + getKind());
 		}
diff --git a/src/core/lombok/javac/JavacTransformer.java b/src/core/lombok/javac/JavacTransformer.java
index 0a4f1f73..625fb283 100644
--- a/src/core/lombok/javac/JavacTransformer.java
+++ b/src/core/lombok/javac/JavacTransformer.java
@@ -27,6 +27,7 @@ import java.util.SortedSet;
 import javax.annotation.processing.Messager;
 
 import com.sun.source.util.Trees;
+import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.JCTree.JCAnnotation;
 import com.sun.tools.javac.tree.JCTree.JCClassDecl;
 import com.sun.tools.javac.tree.JCTree.JCCompilationUnit;
@@ -114,5 +115,10 @@ public class JavacTransformer {
 			JCCompilationUnit top = (JCCompilationUnit) annotationNode.top().get();
 			handlers.handleAnnotation(top, annotationNode, annotation, priority);
 		}
+		
+		@Override public void visitAnnotationOnTypeUse(JCTree typeUse, JavacNode annotationNode, JCAnnotation annotation) {
+			JCCompilationUnit top = (JCCompilationUnit) annotationNode.top().get();
+			handlers.handleAnnotation(top, annotationNode, annotation, priority);
+		}
 	}
 }
diff --git a/src/core/lombok/javac/handlers/HandleNonNull.java b/src/core/lombok/javac/handlers/HandleNonNull.java
index 81aa1525..9a81ffff 100644
--- a/src/core/lombok/javac/handlers/HandleNonNull.java
+++ b/src/core/lombok/javac/handlers/HandleNonNull.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2013-2014 The Project Lombok Authors.
+ * Copyright (C) 2013-2019 The Project Lombok Authors.
  * 
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -74,12 +74,24 @@ public class HandleNonNull extends JavacAnnotationHandler<NonNull> {
 			return;
 		}
 		
-		if (annotationNode.up().getKind() != Kind.ARGUMENT) 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) annotationNode.up().up().get();
+			declaration = (JCMethodDecl) paramNode.up().get();
 		} catch (Exception e) {
 			return;
 		}
@@ -93,7 +105,7 @@ public class HandleNonNull extends JavacAnnotationHandler<NonNull> {
 		// 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(), annotationNode.up(), annotationNode), ast, annotationNode.getContext());
+		JCStatement nullCheck = recursiveSetGeneratedBy(generateNullCheck(annotationNode.getTreeMaker(), paramNode, annotationNode), ast, annotationNode.getContext());
 		
 		if (nullCheck == null) {
 			// @NonNull applied to a primitive. Kinda pointless. Let's generate a warning.
@@ -103,7 +115,7 @@ public class HandleNonNull extends JavacAnnotationHandler<NonNull> {
 		
 		List<JCStatement> statements = declaration.body.stats;
 		
-		String expectedName = annotationNode.up().getName();
+		String expectedName = paramNode.getName();
 		
 		/* Abort if the null check is already there, delving into try and synchronized statements */ {
 			List<JCStatement> stats = statements;
diff --git a/src/core/lombok/javac/handlers/JavacHandlerUtil.java b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
index 5fd17388..f08098d2 100644
--- a/src/core/lombok/javac/handlers/JavacHandlerUtil.java
+++ b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
@@ -1462,7 +1462,7 @@ public class JavacHandlerUtil {
 	 * variable name as message.
 	 */
 	public static JCStatement generateNullCheck(JavacTreeMaker maker, JavacNode variable, JavacNode source) {
-		return generateNullCheck(maker, variable, (JCVariableDecl)variable.get(), source);
+		return generateNullCheck(maker, variable, (JCVariableDecl) variable.get(), source);
 	}
 
 	/**