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;
@@ -334,57 +336,10 @@ public void checkTypeParameterNullnessForFunctionReturnType(
    * @param lhsType type for the lhs of the assignment
    * @param rhsType type for the rhs of the assignment
    */
-  private boolean compareNullabilityAnnotations(Type.ClassType lhsType, Type.ClassType 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
-      if (lhsTypeArgument.getTypeArguments().length() > 0) {
-        if (!compareNullabilityAnnotations(
-            (Type.ClassType) lhsTypeArgument, (Type.ClassType) rhsTypeArgument)) {
-          return false;
-        }
-      }
-    }
-    return true;
+  private boolean compareNullabilityAnnotations(Type lhsType, Type rhsType) {
+    // it is fair to assume rhyType should be the same as lhsType as the Java compiler has passed
+    // before NullAway.
+    return lhsType.accept(COMPARE_NULLABILITY_VISITOR, rhsType);
   }
 
   /**
@@ -433,12 +388,16 @@ private Type.ClassType typeWithPreservedAnnotations(ParameterizedTypeTree tree)
       TypeMetadata typeMetadata =
           new TypeMetadata(new TypeMetadata.Annotations(nullableAnnotationCompound));
       Type currentTypeArgType = castToNonNull(ASTHelpers.getType(curTypeArg));
-      if (currentTypeArgType.getTypeArguments().size() > 0) {
+      List<Type> genericArgType = currentTypeArgType.accept(TYPE_ARG_VISITOR, null);
+      if (genericArgType.size() > 0) {
         // nested generic type; recursively preserve its nullability type argument annotations
-        currentTypeArgType = typeWithPreservedAnnotations((ParameterizedTypeTree) curTypeArg);
+        // currentTypeArgType = typeWithPreservedAnnotations((ParameterizedTypeTree) curTypeArg);
+        // visitor based approach
+        currentTypeArgType = curTypeArg.accept(new PreservedAnnotationTreeVisitor(), null);
       }
-      Type.ClassType newTypeArgType =
-          (Type.ClassType) currentTypeArgType.cloneWithMetadata(typeMetadata);
+      // Type.ClassType newTypeArgType = (Type.ClassType)
+      // currentTypeArgType.cloneWithMetadata(typeMetadata);
+      Type newTypeArgType = currentTypeArgType.cloneWithMetadata(typeMetadata);
       newTypeArgs.add(newTypeArgType);
     }
     Type.ClassType finalType =
@@ -547,6 +506,173 @@ public void compareGenericTypeParameterNullabilityForCall(
     }
   }
 
+  /** To dispatch the logic to obtain the generic type arguments for different Types */
+  private 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) {
+          Type.ClassType classArg = (Type.ClassType) t.getComponentType();
+          return classArg.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 */
+  private final Type.Visitor<Boolean, Type> COMPARE_NULLABILITY_VISITOR =
+      new Types.DefaultTypeVisitor<Boolean, Type>() {
+        @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(COMPARE_NULLABILITY_VISITOR, rhsTypeArgument)) {
+                return false;
+              }
+            }
+          }
+          return true;
+        }
+
+        @Override
+        public Boolean visitArrayType(Type.ArrayType lhsType, Type rhsType) {
+          Type.ArrayType arrRhsType = (Type.ArrayType) rhsType;
+          return lhsType
+              .getComponentType()
+              .accept(COMPARE_NULLABILITY_VISITOR, arrRhsType.getComponentType());
+        }
+
+        @Override
+        @Nullable
+        public Boolean visitType(Type t, Type type) {
+          return null;
+        }
+      };
+
+  /** Visitor For Handling Different Tree Types */
+  public class PreservedAnnotationTreeVisitor extends SimpleTreeVisitor<Type, Boolean> {
+    @Override
+    public Type visitArrayType(ArrayTypeTree tree, Boolean p) {
+      ParameterizedTypeTree paramTree = (ParameterizedTypeTree) tree.getType();
+      Type.ClassType classType =
+          (Type.ClassType) paramTree.accept(new PreservedAnnotationTreeVisitor(), null);
+      return new Type.ArrayType(classType, classType.tsym);
+    }
+
+    @Override
+    public Type visitParameterizedType(ParameterizedTypeTree tree, Boolean 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(), 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()