Infer generics for assignments

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

@@ -277,6 +277,8 @@ public Config getConfig() {
    */
   private final Map<ExpressionTree, Nullness> computedNullnessMap = new LinkedHashMap<>();
 
+  private GenericsChecks genericsChecks = new GenericsChecks();
+
   /**
    * Error Prone requires us to have an empty constructor for each Plugin, in addition to the
    * constructor taking an ErrorProneFlags object. This constructor should not be used anywhere
@@ -500,7 +502,7 @@ public Description matchAssignment(AssignmentTree tree, VisitorState state) {
     }
     // generics check
     if (lhsType != null && config.isJSpecifyMode()) {
-      GenericsChecks.checkTypeParameterNullnessForAssignability(tree, this, state);
+      genericsChecks.checkTypeParameterNullnessForAssignability(tree, this, state);
     }
 
     if (config.isJSpecifyMode() && tree.getVariable() instanceof ArrayAccessTree) {
@@ -1494,7 +1496,7 @@ public Description matchVariable(VariableTree tree, VisitorState state) {
     }
     VarSymbol symbol = ASTHelpers.getSymbol(tree);
     if (tree.getInitializer() != null && config.isJSpecifyMode()) {
-      GenericsChecks.checkTypeParameterNullnessForAssignability(tree, this, state);
+      genericsChecks.checkTypeParameterNullnessForAssignability(tree, this, state);
     }
     if (!config.isLegacyAnnotationLocation()) {
       checkNullableAnnotationPositionInType(
@@ -1662,6 +1664,7 @@ public Description matchClass(ClassTree tree, VisitorState state) {
       class2Entities.clear();
       class2ConstructorUninit.clear();
       computedNullnessMap.clear();
+      genericsChecks.clearCache();
       EnclosingEnvironmentNullness.instance(state.context).clear();
     } else if (classAnnotationIntroducesPartialMarking(classSymbol)) {
       // Handle the case where the top-class is unannotated, but there is a @NullMarked annotation
@@ -1880,7 +1883,7 @@ private Description handleInvocation(
               Nullness.paramHasNullableAnnotation(methodSymbol, i, config)
                   ? Nullness.NULLABLE
                   : ((config.isJSpecifyMode() && tree instanceof MethodInvocationTree)
-                      ? GenericsChecks.getGenericParameterNullnessAtInvocation(
+                      ? genericsChecks.getGenericParameterNullnessAtInvocation(
                           i, methodSymbol, (MethodInvocationTree) tree, state, config)
                       : Nullness.NONNULL);
         }

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

@@ -38,6 +38,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.stream.Collectors;
 import javax.lang.model.type.ExecutableType;
 import javax.lang.model.type.TypeKind;
 import javax.lang.model.type.TypeVariable;
@@ -46,8 +47,12 @@
 /** Methods for performing checks related to generic types and nullability. */
 public final class GenericsChecks {
 
-  /** Do not instantiate; all methods should be static */
-  private GenericsChecks() {}
+  /**
+   * Maps a MethodInvocationTree to a set of type variables that are mapped to their inferred types.
+   * Any generic type parameter that are not explicitly stated are inferred and cached in this
+   * field.
+   */
+  private final Map<MethodInvocationTree, Map<TypeVariable, Type>> inferredTypes = new HashMap<>();
 
   /**
    * Checks that for an instantiated generic type, {@code @Nullable} types are only used for type
@@ -413,13 +418,16 @@ private static void reportInvalidOverridingMethodParamTypeError(
    * @param analysis the analysis object
    * @param state the visitor state
    */
-  public static void checkTypeParameterNullnessForAssignability(
+  public void checkTypeParameterNullnessForAssignability(
       Tree tree, NullAway analysis, VisitorState state) {
     Config config = analysis.getConfig();
     if (!config.isJSpecifyMode()) {
       return;
     }
     Type lhsType = getTreeType(tree, config);
+    if (lhsType == null) {
+      return;
+    }
     Tree rhsTree;
     if (tree instanceof VariableTree) {
       VariableTree varTree = (VariableTree) tree;
@@ -428,14 +436,49 @@ public static void checkTypeParameterNullnessForAssignability(
       AssignmentTree assignmentTree = (AssignmentTree) tree;
       rhsTree = assignmentTree.getExpression();
     }
+
+    if (rhsTree instanceof MethodInvocationTree) {
+      MethodInvocationTree methodInvocationTree = (MethodInvocationTree) rhsTree;
+      Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol(methodInvocationTree);
+      // update inferredTypes cache for assignments
+      // generic method call with no explicit generic arguments
+      if (methodSymbol.type instanceof Type.ForAll
+          && methodInvocationTree.getTypeArguments().isEmpty()) {
+        Type returnType = methodSymbol.getReturnType();
+        Map<TypeVariable, Type> genericNullness =
+            returnType.accept(new InferTypeVisitor(config), lhsType);
+        if (genericNullness != null) {
+          inferredTypes.put(methodInvocationTree, genericNullness);
+        }
+      }
+    }
     // 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 rhsType = getTreeType(rhsTree, config);
 
-    if (lhsType != null && rhsType != null) {
+    if (rhsType != null && rhsTree instanceof MethodInvocationTree) {
+      // recreate rhsType using inferred types
+      MethodInvocationTree methodInvocationTree = (MethodInvocationTree) rhsTree;
+      Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol(methodInvocationTree);
+      if (inferredTypes.containsKey(methodInvocationTree)) {
+        Map<TypeVariable, Type> genericNullness = inferredTypes.get(methodInvocationTree);
+        List<Type> keyTypeList =
+            genericNullness.keySet().stream()
+                .map(typeVar -> (Type) typeVar)
+                .collect(Collectors.toList());
+        com.sun.tools.javac.util.List<Type> from = com.sun.tools.javac.util.List.from(keyTypeList);
+        com.sun.tools.javac.util.List<Type> to =
+            com.sun.tools.javac.util.List.from(genericNullness.values());
+        rhsType =
+            TypeSubstitutionUtils.subst(
+                state.getTypes(), methodSymbol.getReturnType(), from, to, config);
+      }
+    }
+
+    if (rhsType != null) {
       boolean isAssignmentValid = subtypeParameterNullability(lhsType, rhsType, state, config);
       if (!isAssignmentValid) {
         reportInvalidAssignmentInstantiationError(tree, lhsType, rhsType, state, analysis);
@@ -929,7 +972,7 @@ private static Type substituteTypeArgsInGenericMethodType(
    * @return Nullness of parameter at {@code paramIndex}, or {@code NONNULL} if the call does not
    *     invoke an instance method
    */
-  public static Nullness getGenericParameterNullnessAtInvocation(
+  public Nullness getGenericParameterNullnessAtInvocation(
       int paramIndex,
       Symbol.MethodSymbol invokedMethodSymbol,
       MethodInvocationTree tree,
@@ -949,6 +992,21 @@ public static Nullness getGenericParameterNullnessAtInvocation(
               getTypeNullness(substitutedParamTypes.get(paramIndex), config), Nullness.NULLABLE)) {
         return Nullness.NULLABLE;
       }
+      // check nullness of inferred types
+      if (inferredTypes.containsKey(tree)) {
+        Map<TypeVariable, Type> genericNullness = inferredTypes.get(tree);
+        List<Symbol.VarSymbol> parameters = invokedMethodSymbol.getParameters();
+        if (genericNullness.containsKey(parameters.get(paramIndex).type)) {
+          Type genericType = parameters.get(paramIndex).type;
+          Type inferredGenericType = genericNullness.get(genericType);
+          if (inferredGenericType != null
+              && Objects.equals(getTypeNullness(inferredGenericType, config), Nullness.NULLABLE)) {
+            return Nullness.NULLABLE;
+          } else {
+            return Nullness.NONNULL;
+          }
+        }
+      }
     }
 
     if (!(tree.getMethodSelect() instanceof MemberSelectTree) || invokedMethodSymbol.isStatic()) {
@@ -1157,6 +1215,10 @@ public static boolean passingLambdaOrMethodRefWithGenericReturnToUnmarkedCode(
     return callingUnannotated;
   }
 
+  public void clearCache() {
+    inferredTypes.clear();
+  }
+
   public static boolean isNullableAnnotated(Type type, Config config) {
     return Nullness.hasNullableAnnotation(type.getAnnotationMirrors().stream(), config);
   }

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

@@ -0,0 +1,70 @@
+package com.uber.nullaway.generics;
+
+import com.sun.tools.javac.code.Type;
+import com.sun.tools.javac.code.Types;
+import com.uber.nullaway.Config;
+import com.uber.nullaway.Nullness;
+import java.util.HashMap;
+import java.util.Map;
+import javax.lang.model.type.TypeVariable;
+import org.jspecify.annotations.Nullable;
+
+/** Visitor that uses two types to infer the type of type variables. */
+public class InferTypeVisitor
+    extends Types.DefaultTypeVisitor<@Nullable Map<TypeVariable, Type>, Type> {
+  private final Config config;
+
+  InferTypeVisitor(Config config) {
+    this.config = config;
+  }
+
+  @Override
+  public @Nullable Map<TypeVariable, Type> visitClassType(Type.ClassType rhsType, Type lhsType) {
+    Map<TypeVariable, Type> genericNullness = new HashMap<>();
+    com.sun.tools.javac.util.List<Type> rhsTypeArguments = rhsType.getTypeArguments();
+    com.sun.tools.javac.util.List<Type> lhsTypeArguments =
+        ((Type.ClassType) lhsType).getTypeArguments();
+    // get the inferred type for each type arguments and add them to genericNullness
+    for (int i = 0; i < rhsTypeArguments.size(); i++) {
+      Type rhsTypeArg = rhsTypeArguments.get(i);
+      Type lhsTypeArg = lhsTypeArguments.get(i);
+      Map<TypeVariable, Type> map = rhsTypeArg.accept(this, lhsTypeArg);
+      if (map != null) {
+        genericNullness.putAll(map);
+      }
+    }
+    return genericNullness.isEmpty() ? null : genericNullness;
+  }
+
+  @Override
+  public @Nullable Map<TypeVariable, Type> visitArrayType(Type.ArrayType rhsType, Type lhsType) {
+    // unwrap the type of the array and call accept on it
+    Type rhsComponentType = rhsType.elemtype;
+    Type lhsComponentType = ((Type.ArrayType) lhsType).elemtype;
+    Map<TypeVariable, Type> genericNullness = rhsComponentType.accept(this, lhsComponentType);
+    return genericNullness;
+  }
+
+  @Override
+  public Map<TypeVariable, Type> visitTypeVar(Type.TypeVar rhsType, Type lhsType) {
+    Map<TypeVariable, Type> genericNullness = new HashMap<>();
+    Boolean isLhsNullable =
+        Nullness.hasNullableAnnotation(lhsType.getAnnotationMirrors().stream(), config);
+    Type upperBound = rhsType.getUpperBound();
+    Boolean isRhsNullable =
+        Nullness.hasNullableAnnotation(upperBound.getAnnotationMirrors().stream(), config);
+    if (!isLhsNullable) { // lhsType is NonNull, we can just use this
+      genericNullness.put(rhsType, lhsType);
+    } else if (isRhsNullable) { // lhsType & rhsType are Nullable, can use lhs for inference
+      genericNullness.put(rhsType, lhsType);
+    } else { // rhs can't be nullable, use upperbound
+      genericNullness.put(rhsType, upperBound);
+    }
+    return genericNullness;
+  }
+
+  @Override
+  public @Nullable Map<TypeVariable, Type> visitType(Type t, Type type) {
+    return null;
+  }
+}