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()) {
+      handleNullabilityOnNestedClass(
+          ((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()) {
+      handleNullabilityOnNestedClass(
+          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()) {
+      handleNullabilityOnNestedClass(
+          tree.getModifiers().getAnnotations(), tree.getType(), tree, state);
+    }
 
     if (symbol.type.isPrimitive() && tree.getInitializer() != null) {
       doUnboxingCheck(state, tree.getInitializer());
@@ -1487,6 +1509,70 @@
     return Description.NO_MATCH;
   }
 
+  private static boolean isOnlyTypeAnnotation(Symbol anno) {
+    Target target = anno.getAnnotation(Target.class);
+    ImmutableSet<ElementType> elementTypes =
+        target == null ? ImmutableSet.of() : ImmutableSet.copyOf(target.value());
+    return elementTypes.contains(TYPE_USE);
+  }
+
+  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 only if annotation is not type-use only
+    return !(elementTypes.equals(ImmutableSet.of(ElementType.TYPE_USE))
+        || TYPE_USE_OR_TYPE_PARAMETER.containsAll(elementTypes));
+  }
+
+  private void handleNullabilityOnNestedClass(
+      List<? extends AnnotationTree> annotations,
+      Tree typeTree,
+      Tree errorReportingTree,
+      VisitorState state) {
+    if (!(typeTree instanceof JCTree.JCFieldAccess)) {
+      return;
+    }
+    JCTree.JCFieldAccess fieldAccess = (JCTree.JCFieldAccess) typeTree;
+
+    int endOfOuterType = state.getEndPosition(fieldAccess.getExpression());
+    int startOfType = ((JCTree) typeTree).getStartPosition();
+
+    for (AnnotationTree annotation : annotations) {
+      Symbol sym = ASTHelpers.getSymbol(annotation);
+      if (sym == null) {
+        continue;
+      }
+      if (!isOnlyTypeAnnotation(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;
+      }
+
+      String qualifiedName = sym.getQualifiedName().toString();
+      if (!isNullableAnnotation(qualifiedName, config)) {
+        continue;
+      }
+
+      if (state.getEndPosition(annotation) < endOfOuterType) {
+
+        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(errorReportingTree), 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));
     }
   }
 
@@ -327,7 +329,8 @@ public static Stream<? extends AnnotationMirror> getTypeUseAnnotations(
    * @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 +343,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.
+    // Outside of JSpecify mode, annotations which are *not* on the inner type are not treated as
+    // being applied to the inner type. This is bypassed when the LegacyAnnotationLocations flag is
+    // passed, in which case 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 +355,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 +376,32 @@ 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);
+    }
+    if (!hasNestedClass(symbol.type)) {
+      if (innerTypeCount > 0) {
+        return true;
+      }
+    }
+    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);
   }
 
   /**