Enforce correct type-use annotation locations for nested types

nullaway/src/main/java/com/uber/nullaway/ErrorMessage.java

@@ -59,6 +59,7 @@ public enum MessageTypes {
     WRONG_OVERRIDE_RETURN_GENERIC,
     WRONG_OVERRIDE_PARAM_GENERIC,
     ASSIGN_NULLABLE_TO_NONNULL_ARRAY,
+    NULLABLE_ON_WRONG_NESTED_CLASS_LEVEL
   }
 
   public String getMessage() {

nullaway/src/main/java/com/uber/nullaway/NullAway.java

@@ -33,6 +33,9 @@
 import static com.uber.nullaway.ErrorBuilder.errMsgForInitializer;
 import static com.uber.nullaway.NullabilityUtil.castToNonNull;
 import static com.uber.nullaway.NullabilityUtil.isArrayElementNullable;
+import static com.uber.nullaway.Nullness.isNullableAnnotation;
+import static java.lang.annotation.ElementType.TYPE_PARAMETER;
+import static java.lang.annotation.ElementType.TYPE_USE;
 
 import com.google.auto.service.AutoService;
 import com.google.auto.value.AutoValue;
@@ -53,6 +56,8 @@
 import com.google.errorprone.matchers.Matchers;
 import com.google.errorprone.suppliers.Suppliers;
 import com.google.errorprone.util.ASTHelpers;
+import com.sun.source.tree.AnnotatedTypeTree;
+import com.sun.source.tree.AnnotationTree;
 import com.sun.source.tree.ArrayAccessTree;
 import com.sun.source.tree.AssignmentTree;
 import com.sun.source.tree.BinaryTree;
@@ -98,6 +103,8 @@
 import com.uber.nullaway.handlers.Handler;
 import com.uber.nullaway.handlers.Handlers;
 import com.uber.nullaway.handlers.MethodAnalysisContext;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Target;
 import java.util.ArrayList;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
@@ -187,6 +194,8 @@
   static final String CORE_CHECK_NAME = "NullAway.<core>";
 
   private static final Matcher<ExpressionTree> THIS_MATCHER = NullAway::isThisIdentifierMatcher;
+  private static final ImmutableSet<ElementType> TYPE_USE_OR_TYPE_PARAMETER =
+      ImmutableSet.of(TYPE_USE, TYPE_PARAMETER);
 
   private final Predicate<MethodInvocationNode> nonAnnotatedMethod;
 
@@ -570,6 +579,11 @@
         || symbol instanceof ModuleElement) {
       return Description.NO_MATCH;
     }
+    if ((tree.getExpression() instanceof AnnotatedTypeTree)
+        && !config.isLegacyAnnotationLocation()) {
+      checkNullableAnnotationPositionInType(
+          ((AnnotatedTypeTree) tree.getExpression()).getAnnotations(), tree, tree, state);
+    }
 
     Description badDeref = matchDereference(tree.getExpression(), tree, state);
     if (!badDeref.equals(Description.NO_MATCH)) {
@@ -638,6 +652,10 @@
     if (!withinAnnotatedCode(state)) {
       return Description.NO_MATCH;
     }
+    if (!config.isLegacyAnnotationLocation()) {
+      checkNullableAnnotationPositionInType(
+          tree.getModifiers().getAnnotations(), tree.getReturnType(), tree, state);
+    }
     // if the method is overriding some other method,
     // check that nullability annotations are consistent with
     // overridden method (if overridden method is in an annotated
@@ -1464,6 +1482,10 @@
     if (tree.getInitializer() != null && config.isJSpecifyMode()) {
       GenericsChecks.checkTypeParameterNullnessForAssignability(tree, this, state);
     }
+    if (!config.isLegacyAnnotationLocation()) {
+      checkNullableAnnotationPositionInType(
+          tree.getModifiers().getAnnotations(), tree.getType(), tree, state);
+    }
 
     if (symbol.type.isPrimitive() && tree.getInitializer() != null) {
       doUnboxingCheck(state, tree.getInitializer());
@@ -1487,6 +1509,87 @@
     return Description.NO_MATCH;
   }
 
+  /**
+   * returns true if {@code anno} is a type use annotation; it may also be a declaration annotation
+   */
+  private static boolean isTypeUseAnnotation(Symbol anno) {
+    Target target = anno.getAnnotation(Target.class);
+    ImmutableSet<ElementType> elementTypes =
+        target == null ? ImmutableSet.of() : ImmutableSet.copyOf(target.value());
+    return elementTypes.contains(TYPE_USE);
+  }
+
+  /**
+   * returns true if {@code anno} is a declaration annotation; it may also be a type use annotation
+   */
+  private static boolean isDeclarationAnnotation(Symbol anno) {
+    Target target = anno.getAnnotation(Target.class);
+    if (target == null) {
+      return true;
+    }
+    ImmutableSet<ElementType> elementTypes = ImmutableSet.copyOf(target.value());
+    // Return true for any annotation that is not exclusively a type-use annotation
+    return !(elementTypes.equals(ImmutableSet.of(ElementType.TYPE_USE))
+        || TYPE_USE_OR_TYPE_PARAMETER.containsAll(elementTypes));
+  }
+
+  /**
+   * Checks whether any {@code @Nullable} annotation is at the right location for nested types.
+   * Raises an error iff the type is a field access expression (for an inner class type), the
+   * annotation is type use, and the annotation is not applied on the innermost type.
+   *
+   * @param annotations The annotations to check
+   * @param type The tree representing the type structure
+   * @param tree The tree context (variable, member select or method)
+   * @param state The visitor state
+   */
+  private void checkNullableAnnotationPositionInType(
+      List<? extends AnnotationTree> annotations, Tree type, Tree tree, VisitorState state) {
+
+    // Early return if the type is not a nested or inner class reference.
+    if (!(type instanceof MemberSelectTree)) {
+      return;
+    }
+    MemberSelectTree fieldAccess = (MemberSelectTree) type;
+
+    // Get the end position of the outer type expression. Any nullable annotation before this
+    // position is considered to be on the outer type, which is incorrect.
+    int endOfOuterType = state.getEndPosition(fieldAccess.getExpression());
+    int startOfType = ((JCTree) type).getStartPosition();
+
+    for (AnnotationTree annotation : annotations) {
+      Symbol sym = ASTHelpers.getSymbol(annotation);
+      if (sym == null) {
+        continue;
+      }
+
+      String qualifiedName = sym.getQualifiedName().toString();
+      if (!isNullableAnnotation(qualifiedName, config)) {
+        continue;
+      }
+
+      if (!isTypeUseAnnotation(sym)) {
+        continue;
+      }
+      // If an annotation is declaration ALSO, we check if it is at the correct location. If it is,
+      // we treat it as declaration and skip the checks.
+      if (isDeclarationAnnotation(sym) && state.getEndPosition(annotation) <= startOfType) {
+        continue;
+      }
+
+      if (state.getEndPosition(annotation) < endOfOuterType) {
+        // annotation is not on the inner-most type
+        ErrorMessage errorMessage =
+            new ErrorMessage(
+                MessageTypes.NULLABLE_ON_WRONG_NESTED_CLASS_LEVEL,
+                "Type-use nullability annotations should be applied on inner class");
+
+        state.reportMatch(
+            errorBuilder.createErrorDescription(errorMessage, buildDescription(tree), state, null));
+      }
+    }
+  }
+
   /**
    * Check if an inner class's annotation means this Compilation Unit is partially annotated.
    *

nullaway/src/main/java/com/uber/nullaway/NullabilityUtil.java

@@ -44,6 +44,7 @@
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.TypeAnnotationPosition;
 import com.sun.tools.javac.code.TypeAnnotationPosition.TypePathEntry;
+import com.sun.tools.javac.code.TypeTag;
 import com.sun.tools.javac.code.Types;
 import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.util.JCDiagnostic;
@@ -293,7 +294,7 @@ public static Stream<? extends AnnotationMirror> getAllAnnotationsForParameter(
                 t ->
                     t.position.type.equals(TargetType.METHOD_FORMAL_PARAMETER)
                         && t.position.parameter_index == paramInd
-                        && NullabilityUtil.isDirectTypeUseAnnotation(t, config)));
+                        && NullabilityUtil.isDirectTypeUseAnnotation(t, symbol, config)));
   }
 
   /**
@@ -308,10 +309,11 @@ public static Stream<? extends AnnotationMirror> getTypeUseAnnotations(
       return rawTypeAttributes.filter(
           (t) ->
               t.position.type.equals(TargetType.METHOD_RETURN)
-                  && isDirectTypeUseAnnotation(t, config));
+                  && isDirectTypeUseAnnotation(t, symbol, config));
     } else {
       // filter for annotations directly on the type
-      return rawTypeAttributes.filter(t -> NullabilityUtil.isDirectTypeUseAnnotation(t, config));
+      return rawTypeAttributes.filter(
+          t -> NullabilityUtil.isDirectTypeUseAnnotation(t, symbol, config));
     }
   }
 
@@ -323,11 +325,13 @@ public static Stream<? extends AnnotationMirror> getTypeUseAnnotations(
    * but {@code List<@Nullable T> lst} is not.
    *
    * @param t the annotation and its position in the type
+   * @param symbol the symbol for the annotated element
    * @param config NullAway configuration
    * @return {@code true} if the annotation should be treated as applying directly to the top-level
    *     type, false otherwise
    */
-  private static boolean isDirectTypeUseAnnotation(Attribute.TypeCompound t, Config config) {
+  private static boolean isDirectTypeUseAnnotation(
+      Attribute.TypeCompound t, Symbol symbol, Config config) {
     // location is a list of TypePathEntry objects, indicating whether the annotation is
     // on an array, inner type, wildcard, or type argument. If it's empty, then the
     // annotation is directly on the type.
@@ -340,6 +344,9 @@ private static boolean isDirectTypeUseAnnotation(Attribute.TypeCompound t, Confi
     // In JSpecify mode and without the LegacyAnnotationLocations flag, annotations on array
     // dimensions are *not* treated as applying to the top-level type, consistent with the JSpecify
     // spec.
+    // Annotations which are *not* on the inner type are not treated as being applied to the inner
+    // type. This can be bypassed the LegacyAnnotationLocations flag, in which
+    // annotations on all locations are treated as applying to the inner type.
     // We don't allow mixing of inner types and array dimensions in the same location
     // (i.e. `Foo.@Nullable Bar []` is meaningless).
     // These aren't correct semantics for type use annotations, but a series of hacky
@@ -349,10 +356,13 @@ private static boolean isDirectTypeUseAnnotation(Attribute.TypeCompound t, Confi
     // See https://github.com/uber/NullAway/issues/708
     boolean locationHasInnerTypes = false;
     boolean locationHasArray = false;
+    int innerTypeCount = 0;
+    int nestingDepth = getNestingDepth(symbol.type);
     for (TypePathEntry entry : t.position.location) {
       switch (entry.tag) {
         case INNER_TYPE:
           locationHasInnerTypes = true;
+          innerTypeCount++;
           break;
         case ARRAY:
           if (config.isJSpecifyMode() || !config.isLegacyAnnotationLocation()) {
@@ -367,8 +377,35 @@ private static boolean isDirectTypeUseAnnotation(Attribute.TypeCompound t, Confi
           return false;
       }
     }
-    // Make sure it's not a mix of inner types and arrays for this annotation's location
-    return !(locationHasInnerTypes && locationHasArray);
+    if (config.isLegacyAnnotationLocation()) {
+      // Make sure it's not a mix of inner types and arrays for this annotation's location
+      return !(locationHasInnerTypes && locationHasArray);
+    }
+    // For non-nested classes if there are any inner types in the annotation location the annotation
+    // is considered valid to allow annotations on type arguments.
+    if (!hasNestedClass(symbol.type)) {
+      if (innerTypeCount > 0) {
+        return true;
+      }
+    }
+    // For nested classes the annotation is only valid if it is on the innermost type.
+    return innerTypeCount == nestingDepth - 1;
+  }
+
+  private static int getNestingDepth(Type type) {
+    int depth = 0;
+    for (Type curr = type;
+        curr != null && !curr.hasTag(TypeTag.NONE);
+        curr = curr.getEnclosingType()) {
+      depth++;
+    }
+    return depth;
+  }
+
+  private static boolean hasNestedClass(Type type) {
+    return type.tsym != null
+        && type.tsym.owner instanceof Symbol.ClassSymbol
+        && !(type.tsym.owner.owner instanceof Symbol.PackageSymbol);
   }
 
   /**