JSpecify subtyping checks for arrays (#956)

Add checks for proper array subtyping (including nullability of array
contents) at pseudo-assignments and for method overriding. We do not yet
handle multi-dimensional arrays. Note that JSpecify allows for covariant
array subtyping with respect to nullability (see
jspecify/jspecify#65) which complicates the
implementation a bit. Also some hacks are required since javac does not
propagate the array contents nullability annotation into the types of
the relevant trees in all cases.

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

@@ -495,7 +495,7 @@ public Description matchAssignment(AssignmentTree tree, VisitorState state) {
       doUnboxingCheck(state, tree.getExpression());
     }
     // generics check
-    if (lhsType != null && lhsType.getTypeArguments().length() > 0 && config.isJSpecifyMode()) {
+    if (lhsType != null && config.isJSpecifyMode()) {
       GenericsChecks.checkTypeParameterNullnessForAssignability(tree, this, state);
     }
 

nullaway/src/main/java/com/uber/nullaway/generics/CompareNullabilityVisitor.java → nullaway/src/main/java/com/uber/nullaway/generics/CheckIdenticalNullabilityVisitor.java

@@ -1,22 +1,20 @@
 package com.uber.nullaway.generics;
 
 import com.google.errorprone.VisitorState;
-import com.google.errorprone.util.ASTHelpers;
-import com.sun.tools.javac.code.Attribute;
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.Types;
 import java.util.List;
 import javax.lang.model.type.NullType;
 
 /**
- * Visitor that checks equality of nullability annotations for all nested generic type arguments
- * within a type. Compares the Type it is called upon, i.e. the LHS type and the Type passed as an
- * argument, i.e. The RHS type.
+ * Visitor that checks for identical nullability annotations at all nesting levels within two types.
+ * Compares the Type it is called upon, i.e. the LHS type and the Type passed as an argument, i.e.
+ * The RHS type.
  */
-public class CompareNullabilityVisitor extends Types.DefaultTypeVisitor<Boolean, Type> {
+public class CheckIdenticalNullabilityVisitor extends Types.DefaultTypeVisitor<Boolean, Type> {
   private final VisitorState state;
 
-  CompareNullabilityVisitor(VisitorState state) {
+  CheckIdenticalNullabilityVisitor(VisitorState state) {
     this.state = state;
   }
 
@@ -45,26 +43,8 @@ public Boolean visitClassType(Type.ClassType lhsType, Type rhsType) {
     for (int i = 0; i < lhsTypeArguments.size(); i++) {
       Type lhsTypeArgument = lhsTypeArguments.get(i);
       Type rhsTypeArgument = rhsTypeArguments.get(i);
-      boolean isLHSNullableAnnotated = false;
-      List<Attribute.TypeCompound> lhsAnnotations = lhsTypeArgument.getAnnotationMirrors();
-      // To ensure that we are checking only jspecify nullable annotations
-      Type jspecifyNullableType = GenericsChecks.JSPECIFY_NULLABLE_TYPE_SUPPLIER.get(state);
-      for (Attribute.TypeCompound annotation : lhsAnnotations) {
-        if (ASTHelpers.isSameType(
-            (Type) annotation.getAnnotationType(), jspecifyNullableType, state)) {
-          isLHSNullableAnnotated = true;
-          break;
-        }
-      }
-      boolean isRHSNullableAnnotated = false;
-      List<Attribute.TypeCompound> rhsAnnotations = rhsTypeArgument.getAnnotationMirrors();
-      for (Attribute.TypeCompound annotation : rhsAnnotations) {
-        if (ASTHelpers.isSameType(
-            (Type) annotation.getAnnotationType(), jspecifyNullableType, state)) {
-          isRHSNullableAnnotated = true;
-          break;
-        }
-      }
+      boolean isLHSNullableAnnotated = GenericsChecks.isNullableAnnotated(lhsTypeArgument, state);
+      boolean isRHSNullableAnnotated = GenericsChecks.isNullableAnnotated(rhsTypeArgument, state);
       if (isLHSNullableAnnotated != isRHSNullableAnnotated) {
         return false;
       }
@@ -85,7 +65,14 @@ public Boolean visitArrayType(Type.ArrayType lhsType, Type rhsType) {
       return true;
     }
     Type.ArrayType arrRhsType = (Type.ArrayType) rhsType;
-    return lhsType.getComponentType().accept(this, arrRhsType.getComponentType());
+    Type lhsComponentType = lhsType.getComponentType();
+    Type rhsComponentType = arrRhsType.getComponentType();
+    boolean isLHSNullableAnnotated = GenericsChecks.isNullableAnnotated(lhsComponentType, state);
+    boolean isRHSNullableAnnotated = GenericsChecks.isNullableAnnotated(rhsComponentType, state);
+    if (isRHSNullableAnnotated != isLHSNullableAnnotated) {
+      return false;
+    }
+    return lhsComponentType.accept(this, rhsComponentType);
   }
 
   @Override

nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java

@@ -14,6 +14,7 @@
 import com.sun.source.tree.MemberSelectTree;
 import com.sun.source.tree.MethodInvocationTree;
 import com.sun.source.tree.MethodTree;
+import com.sun.source.tree.NewArrayTree;
 import com.sun.source.tree.NewClassTree;
 import com.sun.source.tree.ParameterizedTypeTree;
 import com.sun.source.tree.Tree;
@@ -34,6 +35,7 @@
 import java.util.Map;
 import javax.annotation.Nullable;
 import javax.lang.model.type.ExecutableType;
+import javax.lang.model.type.TypeKind;
 import javax.lang.model.type.TypeVariable;
 
 /** Methods for performing checks related to generic types and nullability. */
@@ -244,7 +246,7 @@ private static void reportInvalidOverridingMethodParamTypeError(
    *
    * <p>This method is required because in some cases, the type returned by {@link
    * com.google.errorprone.util.ASTHelpers#getType(Tree)} fails to preserve type use annotations,
-   * particularly when dealing with {@link com.sun.source.tree.NewClassTree} (e.g., {@code new
+   * e.g., when dealing with {@link com.sun.source.tree.NewClassTree} (e.g., {@code new
    * Foo<@Nullable A>}).
    *
    * @param tree A tree for which we need the type with preserved annotations.
@@ -263,8 +265,29 @@ private static Type getTreeType(Tree tree, VisitorState state) {
         return null;
       }
       return typeWithPreservedAnnotations(paramTypedTree, state);
+    } else if (tree instanceof NewArrayTree
+        && ((NewArrayTree) tree).getType() instanceof AnnotatedTypeTree) {
+      return typeWithPreservedAnnotations(tree, state);
     } else {
-      Type result = ASTHelpers.getType(tree);
+      Type result;
+      if (tree instanceof VariableTree) {
+        // type on the tree itself can be missing nested annotations for arrays; get the type from
+        // the symbol for the declared variable instead
+        VariableTree varTree = (VariableTree) tree;
+        result = ASTHelpers.getSymbol(varTree).type;
+      } else if (tree instanceof AssignmentTree) {
+        // type on the tree itself can be missing nested annotations for arrays; get the type from
+        // the symbol for the assigned location instead, if available
+        AssignmentTree assignmentTree = (AssignmentTree) tree;
+        Symbol lhsSymbol = ASTHelpers.getSymbol(assignmentTree.getVariable());
+        if (lhsSymbol != null) {
+          result = lhsSymbol.type;
+        } else {
+          result = ASTHelpers.getType(assignmentTree);
+        }
+      } else {
+        result = ASTHelpers.getType(tree);
+      }
       if (result != null && result.isRaw()) {
         // bail out of any checking involving raw types for now
         return null;
@@ -289,27 +312,24 @@ public static void checkTypeParameterNullnessForAssignability(
     if (!analysis.getConfig().isJSpecifyMode()) {
       return;
     }
-    Tree lhsTree;
+    Type lhsType = getTreeType(tree, state);
     Tree rhsTree;
     if (tree instanceof VariableTree) {
       VariableTree varTree = (VariableTree) tree;
-      lhsTree = varTree.getType();
       rhsTree = varTree.getInitializer();
     } else {
       AssignmentTree assignmentTree = (AssignmentTree) tree;
-      lhsTree = assignmentTree.getVariable();
       rhsTree = assignmentTree.getExpression();
     }
     // rhsTree can be null for a VariableTree.  Also, we don't need to do a check
     // if rhsTree is the null literal
     if (rhsTree == null || rhsTree.getKind().equals(Tree.Kind.NULL_LITERAL)) {
       return;
     }
-    Type lhsType = getTreeType(lhsTree, state);
     Type rhsType = getTreeType(rhsTree, state);
 
     if (lhsType != null && rhsType != null) {
-      boolean isAssignmentValid = compareNullabilityAnnotations(lhsType, rhsType, state);
+      boolean isAssignmentValid = subtypeParameterNullability(lhsType, rhsType, state);
       if (!isAssignmentValid) {
         reportInvalidAssignmentInstantiationError(tree, lhsType, rhsType, state, analysis);
       }
@@ -342,7 +362,7 @@ public static void checkTypeParameterNullnessForFunctionReturnType(
     Type returnExpressionType = getTreeType(retExpr, state);
     if (formalReturnType != null && returnExpressionType != null) {
       boolean isReturnTypeValid =
-          compareNullabilityAnnotations(formalReturnType, returnExpressionType, state);
+          subtypeParameterNullability(formalReturnType, returnExpressionType, state);
       if (!isReturnTypeValid) {
         reportInvalidReturnTypeError(
             retExpr, formalReturnType, returnExpressionType, state, analysis);
@@ -351,9 +371,8 @@ public static void checkTypeParameterNullnessForFunctionReturnType(
   }
 
   /**
-   * Compare two types from an assignment for identical type parameter nullability, recursively
-   * checking nested generic types. See <a
-   * href="https://jspecify.dev/docs/spec/#nullness-delegating-subtyping">the JSpecify
+   * Compare two types for identical type parameter nullability, recursively checking nested generic
+   * types. See <a href="https://jspecify.dev/docs/spec/#nullness-delegating-subtyping">the JSpecify
    * specification</a> and <a
    * href="https://docs.oracle.com/javase/specs/jls/se14/html/jls-4.html#jls-4.10.2">the JLS
    * subtyping rules for class and interface types</a>.
@@ -362,11 +381,39 @@ public static void checkTypeParameterNullnessForFunctionReturnType(
    * @param rhsType type for the rhs of the assignment
    * @param state the visitor state
    */
-  private static boolean compareNullabilityAnnotations(
+  private static boolean identicalTypeParameterNullability(
       Type lhsType, Type rhsType, VisitorState state) {
-    // it is fair to assume rhyType should be the same as lhsType as the Java compiler has passed
-    // before NullAway.
-    return lhsType.accept(new CompareNullabilityVisitor(state), rhsType);
+    return lhsType.accept(new CheckIdenticalNullabilityVisitor(state), rhsType);
+  }
+
+  /**
+   * Like {@link #identicalTypeParameterNullability(Type, Type, VisitorState)}, but allows for
+   * covariant array subtyping at the top level.
+   *
+   * @param lhsType type for the lhs of the assignment
+   * @param rhsType type for the rhs of the assignment
+   * @param state the visitor state
+   */
+  private static boolean subtypeParameterNullability(
+      Type lhsType, Type rhsType, VisitorState state) {
+    if (lhsType.getKind().equals(TypeKind.ARRAY) && rhsType.getKind().equals(TypeKind.ARRAY)) {
+      // for array types we must allow covariance, i.e., an array of @NonNull references is a
+      // subtype of an array of @Nullable references; see
+      // https://github.com/jspecify/jspecify/issues/65
+      Type.ArrayType lhsArrayType = (Type.ArrayType) lhsType;
+      Type.ArrayType rhsArrayType = (Type.ArrayType) rhsType;
+      Type lhsComponentType = lhsArrayType.getComponentType();
+      Type rhsComponentType = rhsArrayType.getComponentType();
+      boolean isLHSNullableAnnotated = isNullableAnnotated(lhsComponentType, state);
+      boolean isRHSNullableAnnotated = isNullableAnnotated(rhsComponentType, state);
+      // an array of @Nullable references is _not_ a subtype of an array of @NonNull references
+      if (isRHSNullableAnnotated && !isLHSNullableAnnotated) {
+        return false;
+      }
+      return identicalTypeParameterNullability(lhsComponentType, rhsComponentType, state);
+    } else {
+      return identicalTypeParameterNullability(lhsType, rhsType, state);
+    }
   }
 
   /**
@@ -378,9 +425,8 @@ private static boolean compareNullabilityAnnotations(
    * @param state the visitor state
    * @return A Type with preserved annotations.
    */
-  private static Type.ClassType typeWithPreservedAnnotations(
-      ParameterizedTypeTree tree, VisitorState state) {
-    return (Type.ClassType) tree.accept(new PreservedAnnotationTreeVisitor(state), null);
+  private static Type typeWithPreservedAnnotations(Tree tree, VisitorState state) {
+    return tree.accept(new PreservedAnnotationTreeVisitor(state), null);
   }
 
   /**
@@ -407,28 +453,40 @@ public static void checkTypeParameterNullnessForConditionalExpression(
     Tree truePartTree = tree.getTrueExpression();
     Tree falsePartTree = tree.getFalseExpression();
 
-    Type condExprType = getTreeType(tree, state);
+    Type condExprType = getConditionalExpressionType(tree, state);
     Type truePartType = getTreeType(truePartTree, state);
     Type falsePartType = getTreeType(falsePartTree, state);
     // The condExpr type should be the least-upper bound of the true and false part types.  To check
     // the nullability annotations, we check that the true and false parts are assignable to the
     // type of the whole expression
     if (condExprType != null) {
       if (truePartType != null) {
-        if (!compareNullabilityAnnotations(condExprType, truePartType, state)) {
+        if (!subtypeParameterNullability(condExprType, truePartType, state)) {
           reportMismatchedTypeForTernaryOperator(
               truePartTree, condExprType, truePartType, state, analysis);
         }
       }
       if (falsePartType != null) {
-        if (!compareNullabilityAnnotations(condExprType, falsePartType, state)) {
+        if (!subtypeParameterNullability(condExprType, falsePartType, state)) {
           reportMismatchedTypeForTernaryOperator(
               falsePartTree, condExprType, falsePartType, state, analysis);
         }
       }
     }
   }
 
+  @Nullable
+  private static Type getConditionalExpressionType(
+      ConditionalExpressionTree tree, VisitorState state) {
+    // hack: sometimes array nullability doesn't get computed correctly for a conditional expression
+    // on the RHS of an assignment.  So, look at the type of the assignment tree.
+    Tree parent = state.getPath().getParentPath().getLeaf();
+    if (parent instanceof AssignmentTree || parent instanceof VariableTree) {
+      return getTreeType(parent, state);
+    }
+    return getTreeType(tree, state);
+  }
+
   /**
    * Checks that for each parameter p at a call, the type parameter nullability for p's type matches
    * that of the corresponding formal parameter. If a mismatch is found, report an error.
@@ -459,7 +517,7 @@ public static void compareGenericTypeParameterNullabilityForCall(
       if (!formalParameter.getTypeArguments().isEmpty()) {
         Type actualParameter = getTreeType(actualParams.get(i), state);
         if (actualParameter != null) {
-          if (!compareNullabilityAnnotations(formalParameter, actualParameter, state)) {
+          if (!subtypeParameterNullability(formalParameter, actualParameter, state)) {
             reportInvalidParametersNullabilityError(
                 formalParameter, actualParameter, actualParams.get(i), state, analysis);
           }
@@ -474,7 +532,7 @@ public static void compareGenericTypeParameterNullabilityForCall(
         for (int i = formalParams.size() - 1; i < actualParams.size(); i++) {
           Type actualParameter = getTreeType(actualParams.get(i), state);
           if (actualParameter != null) {
-            if (!compareNullabilityAnnotations(varargsElementType, actualParameter, state)) {
+            if (!subtypeParameterNullability(varargsElementType, actualParameter, state)) {
               reportInvalidParametersNullabilityError(
                   varargsElementType, actualParameter, actualParams.get(i), state, analysis);
             }
@@ -781,8 +839,9 @@ private static void checkTypeParameterNullnessForOverridingMethodParameterType(
       Type overridingMethodParameterType = getTreeType(methodParameters.get(i), state);
       Type overriddenMethodParameterType = overriddenMethodParameterTypes.get(i);
       if (overriddenMethodParameterType != null && overridingMethodParameterType != null) {
-        if (!compareNullabilityAnnotations(
-            overriddenMethodParameterType, overridingMethodParameterType, state)) {
+        // allow contravariant subtyping
+        if (!subtypeParameterNullability(
+            overridingMethodParameterType, overriddenMethodParameterType, state)) {
           reportInvalidOverridingMethodParamTypeError(
               methodParameters.get(i),
               overriddenMethodParameterType,
@@ -807,11 +866,14 @@ private static void checkTypeParameterNullnessForOverridingMethodParameterType(
   private static void checkTypeParameterNullnessForOverridingMethodReturnType(
       MethodTree tree, Type overriddenMethodType, NullAway analysis, VisitorState state) {
     Type overriddenMethodReturnType = overriddenMethodType.getReturnType();
-    Type overridingMethodReturnType = getTreeType(tree.getReturnType(), state);
-    if (overriddenMethodReturnType == null || overridingMethodReturnType == null) {
+    // We get the return type from the Symbol; the type attached to tree may not have correct
+    // annotations for array types
+    Type overridingMethodReturnType = ASTHelpers.getSymbol(tree).getReturnType();
+    if (overriddenMethodReturnType.isRaw() || overridingMethodReturnType.isRaw()) {
       return;
     }
-    if (!compareNullabilityAnnotations(
+    // allow covariant subtyping
+    if (!subtypeParameterNullability(
         overriddenMethodReturnType, overridingMethodReturnType, state)) {
       reportInvalidOverridingMethodReturnTypeError(
           tree, overriddenMethodReturnType, overridingMethodReturnType, analysis, state);
@@ -882,4 +944,19 @@ public static boolean passingLambdaOrMethodRefWithGenericReturnToUnmarkedCode(
     }
     return callingUnannotated;
   }
+
+  public static boolean isNullableAnnotated(Type type, VisitorState state) {
+    boolean result = false;
+    // To ensure that we are checking only jspecify nullable annotations
+    Type jspecifyNullableType = JSPECIFY_NULLABLE_TYPE_SUPPLIER.get(state);
+    List<Attribute.TypeCompound> lhsAnnotations = type.getAnnotationMirrors();
+    for (Attribute.TypeCompound annotation : lhsAnnotations) {
+      if (ASTHelpers.isSameType(
+          (Type) annotation.getAnnotationType(), jspecifyNullableType, state)) {
+        result = true;
+        break;
+      }
+    }
+    return result;
+  }
 }

nullaway/src/main/java/com/uber/nullaway/generics/PreservedAnnotationTreeVisitor.java

@@ -7,6 +7,7 @@
 import com.sun.source.tree.AnnotatedTypeTree;
 import com.sun.source.tree.AnnotationTree;
 import com.sun.source.tree.ArrayTypeTree;
+import com.sun.source.tree.NewArrayTree;
 import com.sun.source.tree.ParameterizedTypeTree;
 import com.sun.source.tree.Tree;
 import com.sun.source.util.SimpleTreeVisitor;
@@ -36,6 +37,12 @@ public class PreservedAnnotationTreeVisitor extends SimpleTreeVisitor<Type, Void
     this.state = state;
   }
 
+  @Override
+  public Type visitNewArray(NewArrayTree tree, Void p) {
+    Type elemType = tree.getType().accept(this, null);
+    return new Type.ArrayType(elemType, castToNonNull(ASTHelpers.getType(tree)).tsym);
+  }
+
   @Override
   public Type visitArrayType(ArrayTypeTree tree, Void p) {
     Type elemType = tree.getType().accept(this, null);