Enforce correct type-use annotation locations for nested types (#1045)

Enforce correct type-use annotation locations for nested types as per
JSpecify norms. We enforce type-use annotations to be on the inner class
and raise an error if they are not. For annotations that are both type
use and declaration, we raise an error at an invalid location.

**Current behavior**
```
// All three are ok

@nullable A.B.C foo1 = null;
A.@nullable B.C foo2 = null;
A.B.@nullable C foo3 = null;

```

**New behavior**

```
// BUG: Diagnostic contains: Type-use nullability annotations should be applied on inner class
@nullable A.B.C foo1 = null;
// BUG: Diagnostic contains: Type-use nullability annotations should be applied on inner class
A.@nullable B.C foo2 = null;
A.B.@nullable C foo3 = null;


```

For annotations which are both declaration and annotation and type-use,
only `foo2` throws an error since the location isn't apt for either
scenarios

---------

Co-authored-by: Manu Sridharan <[email protected]>

Author: armughan11

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 @@ public class NullAway extends BugChecker
   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 @@ public Description matchMemberSelect(MemberSelectTree tree, VisitorState state)
         || symbol instanceof ModuleElement) {
       return Description.NO_MATCH;
     }
+    if ((tree.getExpression() instanceof AnnotatedTypeTree)
+        && !config.isLegacyAnnotationLocation()) {
+      checkNullableAnnotationPositionInType(
+          ((AnnotatedTypeTree) tree.getExpression()).getAnnotations(), tree, state);
+    }
 
     Description badDeref = matchDereference(tree.getExpression(), tree, state);
     if (!badDeref.equals(Description.NO_MATCH)) {
@@ -638,6 +652,10 @@ public Description matchMethod(MethodTree tree, VisitorState state) {
     if (!withinAnnotatedCode(state)) {
       return Description.NO_MATCH;
     }
+    if (!config.isLegacyAnnotationLocation()) {
+      checkNullableAnnotationPositionInType(
+          tree.getModifiers().getAnnotations(), tree.getReturnType(), 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 @@ public Description matchVariable(VariableTree tree, VisitorState state) {
     if (tree.getInitializer() != null && config.isJSpecifyMode()) {
       GenericsChecks.checkTypeParameterNullnessForAssignability(tree, this, state);
     }
+    if (!config.isLegacyAnnotationLocation()) {
+      checkNullableAnnotationPositionInType(
+          tree.getModifiers().getAnnotations(), tree.getType(), state);
+    }
 
     if (symbol.type.isPrimitive() && tree.getInitializer() != null) {
       doUnboxingCheck(state, tree.getInitializer());
@@ -1487,6 +1509,85 @@ public Description matchVariable(VariableTree tree, VisitorState state) {
     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 state The visitor state
+   */
+  private void checkNullableAnnotationPositionInType(
+      List<? extends AnnotationTree> annotations, Tree type, VisitorState state) {
+
+    // Early return if the type is not a nested or inner class reference.
+    if (!(type instanceof MemberSelectTree)) {
+      return;
+    }
+
+    // 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(((MemberSelectTree) type).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(type), 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,30 @@ 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 annotations apply to the innermost type.
+    if (!isTypeOfNestedClass(symbol.type)) {
+      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 isTypeOfNestedClass(Type type) {
+    return type.tsym != null && type.tsym.owner instanceof Symbol.ClassSymbol;
   }
 
   /**