Issue #740: Adding visitors for handling different types along with ClassType in Generic Type invariance check

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

@@ -10,6 +10,7 @@
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.AnnotatedTypeTree;
 import com.sun.source.tree.AnnotationTree;
+import com.sun.source.tree.ArrayTypeTree;
 import com.sun.source.tree.AssignmentTree;
 import com.sun.source.tree.ConditionalExpressionTree;
 import com.sun.source.tree.ExpressionTree;
@@ -20,6 +21,7 @@
 import com.sun.source.tree.ParameterizedTypeTree;
 import com.sun.source.tree.Tree;
 import com.sun.source.tree.VariableTree;
+import com.sun.source.util.SimpleTreeVisitor;
 import com.sun.tools.javac.code.Attribute;
 import com.sun.tools.javac.code.BoundKind;
 import com.sun.tools.javac.code.Symbol;
@@ -351,60 +353,12 @@ public static void checkTypeParameterNullnessForFunctionReturnType(
    *
    * @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 compareNullabilityAnnotations(
-      Type.ClassType lhsType, Type.ClassType rhsType, VisitorState state) {
-    Types types = state.getTypes();
-    // The base type of rhsType may be a subtype of lhsType's base type.  In such cases, we must
-    // compare lhsType against the supertype of rhsType with a matching base type.
-    rhsType = (Type.ClassType) types.asSuper(rhsType, lhsType.tsym);
-    // This is impossible, considering the fact that standard Java subtyping succeeds before running
-    // NullAway
-    if (rhsType == null) {
-      throw new RuntimeException("Did not find supertype of " + rhsType + " matching " + lhsType);
-    }
-    List<Type> lhsTypeArguments = lhsType.getTypeArguments();
-    List<Type> rhsTypeArguments = rhsType.getTypeArguments();
-    // This is impossible, considering the fact that standard Java subtyping succeeds before running
-    // NullAway
-    if (lhsTypeArguments.size() != rhsTypeArguments.size()) {
-      throw new RuntimeException(
-          "Number of types arguments in " + rhsType + " does not match " + lhsType);
-    }
-    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
-      for (Attribute.TypeCompound annotation : lhsAnnotations) {
-        if (annotation.getAnnotationType().toString().equals(NULLABLE_NAME)) {
-          isLHSNullableAnnotated = true;
-          break;
-        }
-      }
-      boolean isRHSNullableAnnotated = false;
-      List<Attribute.TypeCompound> rhsAnnotations = rhsTypeArgument.getAnnotationMirrors();
-      // To ensure that we are checking only jspecify nullable annotations
-      for (Attribute.TypeCompound annotation : rhsAnnotations) {
-        if (annotation.getAnnotationType().toString().equals(NULLABLE_NAME)) {
-          isRHSNullableAnnotated = true;
-          break;
-        }
-      }
-      if (isLHSNullableAnnotated != isRHSNullableAnnotated) {
-        return false;
-      }
-      // nested generics
-      if (lhsTypeArgument.getTypeArguments().length() > 0) {
-        if (!compareNullabilityAnnotations(
-            (Type.ClassType) lhsTypeArgument, (Type.ClassType) rhsTypeArgument, state)) {
-          return false;
-        }
-      }
-    }
-    return true;
+      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);
   }
 
   /**
@@ -455,13 +409,11 @@ private static Type.ClassType typeWithPreservedAnnotations(
       TypeMetadata typeMetadata =
           new TypeMetadata(new TypeMetadata.Annotations(nullableAnnotationCompound));
       Type currentTypeArgType = castToNonNull(ASTHelpers.getType(curTypeArg));
-      if (currentTypeArgType.getTypeArguments().size() > 0) {
-        // nested generic type; recursively preserve its nullability type argument annotations
-        currentTypeArgType =
-            typeWithPreservedAnnotations((ParameterizedTypeTree) curTypeArg, state);
+      List<Type> genericArgType = currentTypeArgType.accept(TYPE_ARG_VISITOR, null);
+      if (genericArgType.size() > 0) {
+        currentTypeArgType = curTypeArg.accept(new PreservedAnnotationTreeVisitor(state), null);
       }
-      Type.ClassType newTypeArgType =
-          (Type.ClassType) currentTypeArgType.cloneWithMetadata(typeMetadata);
+      Type newTypeArgType = currentTypeArgType.cloneWithMetadata(typeMetadata);
       newTypeArgs.add(newTypeArgType);
     }
     Type.ClassType finalType =
@@ -577,6 +529,182 @@ public static void compareGenericTypeParameterNullabilityForCall(
     }
   }
 
+  /** To dispatch the logic to obtain the generic type arguments for different Types */
+  private static final Type.Visitor<List<Type>, Void> TYPE_ARG_VISITOR =
+      new Types.DefaultTypeVisitor<List<Type>, Void>() {
+        @Override
+        public List<Type> visitClassType(Type.ClassType t, Void s) {
+          return t.getTypeArguments();
+        }
+
+        @Override
+        public List<Type> visitArrayType(Type.ArrayType t, Void unused) {
+          return t.getComponentType().getTypeArguments();
+        }
+
+        @Override
+        public List<Type> visitCapturedType(Type.CapturedType t, Void s) {
+          return t.wildcard.accept(this, null);
+        }
+
+        @Override
+        public List<Type> visitType(Type t, Void unused) {
+          return Collections.emptyList();
+        }
+      };
+
+  /** To dispatch the logic to0 obtain the generic type arguments for different Types */
+  public static class CompareNullabilityVisitor extends Types.DefaultTypeVisitor<Boolean, Type> {
+    private final VisitorState state;
+
+    CompareNullabilityVisitor(VisitorState state) {
+      this.state = state;
+    }
+
+    @Override
+    public Boolean visitClassType(Type.ClassType lhsType, Type rhsType) {
+      Types types = state.getTypes();
+      // The base type of rhsType may be a subtype of lhsType's base type.  In such cases, we
+      // must
+      // compare lhsType against the supertype of rhsType with a matching base type.
+      rhsType = (Type.ClassType) types.asSuper(rhsType, lhsType.tsym);
+      // This is impossible, considering the fact that standard Java subtyping succeeds before
+      // running
+      // NullAway
+      if (rhsType == null) {
+        throw new RuntimeException("Did not find supertype of " + rhsType + " matching " + lhsType);
+      }
+      List<Type> lhsTypeArguments = lhsType.getTypeArguments();
+      List<Type> rhsTypeArguments = rhsType.getTypeArguments();
+      // This is impossible, considering the fact that standard Java subtyping succeeds before
+      // running
+      // NullAway
+      if (lhsTypeArguments.size() != rhsTypeArguments.size()) {
+        throw new RuntimeException(
+            "Number of types arguments in " + rhsType + " does not match " + lhsType);
+      }
+      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
+        for (Attribute.TypeCompound annotation : lhsAnnotations) {
+          if (annotation.getAnnotationType().toString().equals(NULLABLE_NAME)) {
+            isLHSNullableAnnotated = true;
+            break;
+          }
+        }
+        boolean isRHSNullableAnnotated = false;
+        List<Attribute.TypeCompound> rhsAnnotations = rhsTypeArgument.getAnnotationMirrors();
+        // To ensure that we are checking only jspecify nullable annotations
+        for (Attribute.TypeCompound annotation : rhsAnnotations) {
+          if (annotation.getAnnotationType().toString().equals(NULLABLE_NAME)) {
+            isRHSNullableAnnotated = true;
+            break;
+          }
+        }
+        if (isLHSNullableAnnotated != isRHSNullableAnnotated) {
+          return false;
+        }
+        // nested generics
+        List<Type> genericArgs = lhsTypeArgument.accept(TYPE_ARG_VISITOR, null);
+        if (genericArgs.size() > 0) {
+          if (!lhsTypeArgument.accept(new CompareNullabilityVisitor(state), rhsTypeArgument)) {
+            return false;
+          }
+        }
+      }
+      return true;
+    }
+
+    @Override
+    public Boolean visitArrayType(Type.ArrayType lhsType, Type rhsType) {
+      Type.ArrayType arrRhsType = (Type.ArrayType) rhsType;
+      return lhsType
+          .getComponentType()
+          .accept(new CompareNullabilityVisitor(state), arrRhsType.getComponentType());
+    }
+
+    @Override
+    public Boolean visitType(Type t, Type type) {
+      return true;
+    }
+  }
+
+  /** Visitor For Handling Different Tree Types */
+  public static class PreservedAnnotationTreeVisitor extends SimpleTreeVisitor<Type, Void> {
+
+    private final VisitorState state;
+
+    PreservedAnnotationTreeVisitor(VisitorState state) {
+      this.state = state;
+    }
+
+    @Override
+    public Type visitArrayType(ArrayTypeTree tree, Void p) {
+      ParameterizedTypeTree paramTree = (ParameterizedTypeTree) tree.getType();
+      Type.ClassType classType =
+          (Type.ClassType) paramTree.accept(new PreservedAnnotationTreeVisitor(state), null);
+      return new Type.ArrayType(classType, classType.tsym);
+    }
+
+    @Override
+    public Type visitParameterizedType(ParameterizedTypeTree tree, Void p) {
+      Type.ClassType type = (Type.ClassType) ASTHelpers.getType(tree);
+      Preconditions.checkNotNull(type);
+      Type nullableType = NULLABLE_TYPE_SUPPLIER.get(state);
+      List<? extends Tree> typeArguments = tree.getTypeArguments();
+      List<Type> newTypeArgs = new ArrayList<>();
+      boolean hasNullableAnnotation = false;
+      for (int i = 0; i < typeArguments.size(); i++) {
+        AnnotatedTypeTree annotatedType = null;
+        Tree curTypeArg = typeArguments.get(i);
+        // If the type argument has an annotation, it will either be an AnnotatedTypeTree, or a
+        // ParameterizedTypeTree in the case of a nested generic type
+        if (curTypeArg instanceof AnnotatedTypeTree) {
+          annotatedType = (AnnotatedTypeTree) curTypeArg;
+        } else if (curTypeArg instanceof ParameterizedTypeTree
+            && ((ParameterizedTypeTree) curTypeArg).getType() instanceof AnnotatedTypeTree) {
+          annotatedType = (AnnotatedTypeTree) ((ParameterizedTypeTree) curTypeArg).getType();
+        }
+        List<? extends AnnotationTree> annotations =
+            annotatedType != null ? annotatedType.getAnnotations() : Collections.emptyList();
+        for (AnnotationTree annotation : annotations) {
+          if (ASTHelpers.isSameType(
+              nullableType, ASTHelpers.getType(annotation.getAnnotationType()), state)) {
+            hasNullableAnnotation = true;
+            break;
+          }
+        }
+        // construct a TypeMetadata object containing a nullability annotation if needed
+        com.sun.tools.javac.util.List<Attribute.TypeCompound> nullableAnnotationCompound =
+            hasNullableAnnotation
+                ? com.sun.tools.javac.util.List.from(
+                    Collections.singletonList(
+                        new Attribute.TypeCompound(
+                            nullableType, com.sun.tools.javac.util.List.nil(), null)))
+                : com.sun.tools.javac.util.List.nil();
+        TypeMetadata typeMetadata =
+            new TypeMetadata(new TypeMetadata.Annotations(nullableAnnotationCompound));
+        Type currentTypeArgType = castToNonNull(ASTHelpers.getType(curTypeArg));
+        List<Type> genericArgType = currentTypeArgType.accept(TYPE_ARG_VISITOR, null);
+        if (genericArgType.size() > 0) {
+          // nested generic type; recursively preserve its nullability type argument annotations
+          currentTypeArgType = curTypeArg.accept(new PreservedAnnotationTreeVisitor(state), null);
+        }
+        // Type.ClassType newTypeArgType = (Type.ClassType)
+        // currentTypeArgType.cloneWithMetadata(typeMetadata);
+        Type newTypeArgType = currentTypeArgType.cloneWithMetadata(typeMetadata);
+        newTypeArgs.add(newTypeArgType);
+      }
+      Type.ClassType finalType =
+          new Type.ClassType(
+              type.getEnclosingType(), com.sun.tools.javac.util.List.from(newTypeArgs), type.tsym);
+      return finalType;
+    }
+  }
+
   /**
    * Checks that type parameter nullability is consistent between an overriding method and the
    * corresponding overridden method.

nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java

@@ -760,6 +760,26 @@ public void rawTypes() {
         .doTest();
   }
 
+  @Test
+  public void nestedGenericTypeAssignment() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "class Test {",
+            "  static class A<T extends @Nullable Object> { }",
+            "  static void testPositive() {",
+            "    // BUG: Diagnostic contains: Cannot assign from type",
+            "    A<A<@Nullable String>[]> var = new A<A<String>[]>();",
+            "  }",
+            "  static void testNegative() {",
+            "    A<A<@Nullable String>[]> var = new A<A<@Nullable String>[]>();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   @Test
   public void overrideReturnTypes() {
     makeHelper()