JSpecify: infer generic method type arguments for assignments (#1131)

haewiful · msridhar · web-flow · commit 8786e8847934 · 2025-03-13T12:21:58.000-04:00
This PR infers nullability of generic method type arguments for the case where the result of the call is assigned to a variable, e.g.: ```java class Foo<T extends @nullable Object> { Foo(T t) {} static <U extends @nullable Object> Foo<U> make(U u) { return new Foo<>(u); } void test() { Foo<@nullable Object> f1 = Foo.make(null); // legal Foo<Object> f2 = Foo.make(null); // error reported; arg cannot be @nullable } } ``` See the tests for (many) more examples. --------- Co-authored-by: Manu Sridharan <msridhar@gmail.com>
diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -248,6 +248,11 @@ public Config getConfig() {
    */
   private final Handler handler;
 
+  /** Returns the handler being used for this analysis */
+  public Handler getHandler() {
+    return handler;
+  }
+
   /**
    * entities relevant to field initialization per class. cached for performance. nulled out in
    * {@link #matchClass(ClassTree, VisitorState)}
@@ -277,6 +282,14 @@ public Config getConfig() {
    */
   private final Map<ExpressionTree, Nullness> computedNullnessMap = new LinkedHashMap<>();
 
+  /** Logic and state for generics checking */
+  private final GenericsChecks genericsChecks = new GenericsChecks();
+
+  /** Returns the GenericsChecks object for this analysis, used for generics-related checking */
+  public GenericsChecks getGenericsChecks() {
+    return 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 +513,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 +1507,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 +1675,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
@@ -1885,13 +1899,13 @@ 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);
         }
       }
       if (config.isJSpecifyMode()) {
-        GenericsChecks.compareGenericTypeParameterNullabilityForCall(
+        genericsChecks.compareGenericTypeParameterNullabilityForCall(
             methodSymbol, tree, actualParams, varArgsMethod, this, state);
         if (!methodSymbol.getTypeParameters().isEmpty()) {
           GenericsChecks.checkGenericMethodCallTypeArguments(tree, state, this, config, handler);
@@ -2637,8 +2651,8 @@ private boolean mayBeNullMethodCall(
       return true;
     }
     if (config.isJSpecifyMode()
-        && GenericsChecks.getGenericReturnNullnessAtInvocation(
-                exprSymbol, invocationTree, state, config)
+        && genericsChecks
+            .getGenericReturnNullnessAtInvocation(exprSymbol, invocationTree, state, config)
             .equals(Nullness.NULLABLE)) {
       return true;
     }
@@ -2657,7 +2671,7 @@ public boolean nullnessFromDataflow(VisitorState state, ExpressionTree expr) {
   }
 
   public AccessPathNullnessAnalysis getNullnessAnalysis(VisitorState state) {
-    return AccessPathNullnessAnalysis.instance(state, nonAnnotatedMethod, config, this.handler);
+    return AccessPathNullnessAnalysis.instance(state, nonAnnotatedMethod, this);
   }
 
   private Description matchDereference(
diff --git a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessAnalysis.java b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessAnalysis.java
@@ -28,6 +28,7 @@
 import com.sun.source.util.TreePath;
 import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.Config;
+import com.uber.nullaway.NullAway;
 import com.uber.nullaway.Nullness;
 import com.uber.nullaway.handlers.Handler;
 import com.uber.nullaway.handlers.contract.ContractNullnessStoreInitializer;
@@ -65,10 +66,9 @@ public final class AccessPathNullnessAnalysis {
 
   // Use #instance to instantiate
   private AccessPathNullnessAnalysis(
-      Predicate<MethodInvocationNode> methodReturnsNonNull,
-      VisitorState state,
-      Config config,
-      Handler handler) {
+      Predicate<MethodInvocationNode> methodReturnsNonNull, VisitorState state, NullAway analysis) {
+    Config config = analysis.getConfig();
+    Handler handler = analysis.getHandler();
     apContext =
         AccessPath.AccessPathContext.builder()
             .setImmutableTypes(handler.onRegisterImmutableTypes())
@@ -79,8 +79,7 @@ private AccessPathNullnessAnalysis(
             methodReturnsNonNull,
             state,
             apContext,
-            config,
-            handler,
+            analysis,
             new CoreNullnessStoreInitializer());
     this.dataFlow = new DataFlow(config.assertsEnabled(), handler);
 
@@ -91,8 +90,7 @@ private AccessPathNullnessAnalysis(
               methodReturnsNonNull,
               state,
               apContext,
-              config,
-              handler,
+              analysis,
               new ContractNullnessStoreInitializer());
     }
   }
@@ -103,18 +101,15 @@ private AccessPathNullnessAnalysis(
    * @param state visitor state for the compilation
    * @param methodReturnsNonNull predicate determining whether a method is assumed to return NonNull
    *     value
-   * @param config analysis config
+   * @param analysis instance of NullAway analysis
    * @return instance of the analysis
    */
   public static AccessPathNullnessAnalysis instance(
-      VisitorState state,
-      Predicate<MethodInvocationNode> methodReturnsNonNull,
-      Config config,
-      Handler handler) {
+      VisitorState state, Predicate<MethodInvocationNode> methodReturnsNonNull, NullAway analysis) {
     Context context = state.context;
     AccessPathNullnessAnalysis instance = context.get(FIELD_NULLNESS_ANALYSIS_KEY);
     if (instance == null) {
-      instance = new AccessPathNullnessAnalysis(methodReturnsNonNull, state, config, handler);
+      instance = new AccessPathNullnessAnalysis(methodReturnsNonNull, state, analysis);
       context.put(FIELD_NULLNESS_ANALYSIS_KEY, instance);
     }
     return instance;
diff --git a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java
@@ -37,6 +37,7 @@
 import com.sun.tools.javac.code.TypeTag;
 import com.uber.nullaway.CodeAnnotationInfo;
 import com.uber.nullaway.Config;
+import com.uber.nullaway.NullAway;
 import com.uber.nullaway.NullabilityUtil;
 import com.uber.nullaway.Nullness;
 import com.uber.nullaway.generics.GenericsChecks;
@@ -160,22 +161,24 @@ public class AccessPathNullnessPropagation
 
   private final Handler handler;
 
+  private final GenericsChecks genericsChecks;
+
   private final NullnessStoreInitializer nullnessStoreInitializer;
 
   public AccessPathNullnessPropagation(
       Nullness defaultAssumption,
       Predicate<MethodInvocationNode> methodReturnsNonNull,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
-      Config config,
-      Handler handler,
+      NullAway analysis,
       NullnessStoreInitializer nullnessStoreInitializer) {
     this.defaultAssumption = defaultAssumption;
     this.methodReturnsNonNull = methodReturnsNonNull;
     this.state = state;
     this.apContext = apContext;
-    this.config = config;
-    this.handler = handler;
+    this.config = analysis.getConfig();
+    this.handler = analysis.getHandler();
+    this.genericsChecks = analysis.getGenericsChecks();
     this.nullnessStoreInitializer = nullnessStoreInitializer;
   }
 
@@ -1086,7 +1089,7 @@ private boolean genericReturnIsNullable(MethodInvocationNode node) {
       MethodInvocationTree tree = node.getTree();
       if (tree != null) {
         Nullness nullness =
-            GenericsChecks.getGenericReturnNullnessAtInvocation(
+            genericsChecks.getGenericReturnNullnessAtInvocation(
                 ASTHelpers.getSymbol(tree), tree, state, config);
         return nullness.equals(NULLABLE);
       }
diff --git a/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java b/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java
@@ -26,6 +26,7 @@
 import com.sun.tools.javac.code.TargetType;
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.tree.JCTree;
+import com.sun.tools.javac.util.ListBuffer;
 import com.uber.nullaway.CodeAnnotationInfo;
 import com.uber.nullaway.Config;
 import com.uber.nullaway.ErrorBuilder;
@@ -35,6 +36,7 @@
 import com.uber.nullaway.handlers.Handler;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -46,8 +48,14 @@
 /** 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 representing a call to a generic method to a substitution for its
+   * type arguments. The call must not have any explicit type arguments. The substitution is a map
+   * from type variables for the method to their inferred type arguments (most importantly with
+   * inferred nullability information).
+   */
+  private final Map<MethodInvocationTree, Map<TypeVariable, Type>>
+      inferredSubstitutionsForGenericMethodCalls = new LinkedHashMap<>();
 
   /**
    * Checks that for an instantiated generic type, {@code @Nullable} types are only used for type
@@ -413,13 +421,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;
@@ -435,14 +446,58 @@ public static void checkTypeParameterNullnessForAssignability(
     }
     Type rhsType = getTreeType(rhsTree, config);
 
-    if (lhsType != null && rhsType != null) {
+    if (rhsTree instanceof MethodInvocationTree) {
+      MethodInvocationTree methodInvocationTree = (MethodInvocationTree) rhsTree;
+      Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol(methodInvocationTree);
+      if (methodSymbol.type instanceof Type.ForAll
+          && methodInvocationTree.getTypeArguments().isEmpty()) {
+        // generic method call with no explicit generic arguments
+        // update inferred type arguments based on the assignment context
+        InferSubstitutionViaAssignmentContextVisitor inferVisitor =
+            new InferSubstitutionViaAssignmentContextVisitor(config);
+        Type returnType = methodSymbol.getReturnType();
+        returnType.accept(inferVisitor, lhsType);
+
+        Map<TypeVariable, Type> substitution = inferVisitor.getInferredSubstitution();
+        inferredSubstitutionsForGenericMethodCalls.put(methodInvocationTree, substitution);
+        if (rhsType != null) {
+          // update rhsType with inferred substitution
+          rhsType =
+              substituteInferredTypesForTypeVariables(
+                  state, methodSymbol.getReturnType(), substitution, config);
+        }
+      }
+    }
+
+    if (rhsType != null) {
       boolean isAssignmentValid = subtypeParameterNullability(lhsType, rhsType, state, config);
       if (!isAssignmentValid) {
         reportInvalidAssignmentInstantiationError(tree, lhsType, rhsType, state, analysis);
       }
     }
   }
 
+  /**
+   * Substitutes inferred types for type variables within a type.
+   *
+   * @param state The visitor state
+   * @param targetType The type with type variables on which substitutions will be applied
+   * @param substitution The cache that maps type variables to its inferred types
+   * @param config Configuration for the analysis
+   * @return {@code targetType} with the substitutions applied
+   */
+  private Type substituteInferredTypesForTypeVariables(
+      VisitorState state, Type targetType, Map<TypeVariable, Type> substitution, Config config) {
+    ListBuffer<Type> typeVars = new ListBuffer<>();
+    ListBuffer<Type> inferredTypes = new ListBuffer<>();
+    for (Map.Entry<TypeVariable, Type> entry : substitution.entrySet()) {
+      typeVars.append((Type) entry.getKey());
+      inferredTypes.append(entry.getValue());
+    }
+    return TypeSubstitutionUtils.subst(
+        state.getTypes(), targetType, typeVars.toList(), inferredTypes.toList(), config);
+  }
+
   /**
    * Checks that the nullability of type parameters for a returned expression matches that of the
    * type parameters of the enclosing method's return type.
@@ -613,7 +668,7 @@ public static void checkTypeParameterNullnessForConditionalExpression(
    * @param analysis the analysis object
    * @param state the visitor state
    */
-  public static void compareGenericTypeParameterNullabilityForCall(
+  public void compareGenericTypeParameterNullabilityForCall(
       Symbol.MethodSymbol methodSymbol,
       Tree tree,
       List<? extends ExpressionTree> actualParams,
@@ -640,7 +695,7 @@ public static void compareGenericTypeParameterNullabilityForCall(
             TypeSubstitutionUtils.memberType(state.getTypes(), enclosingType, methodSymbol, config);
       }
     }
-    // substitute type arguments for generic methods
+    // substitute type arguments for generic methods with explicit type arguments
     if (tree instanceof MethodInvocationTree && methodSymbol.type instanceof Type.ForAll) {
       invokedMethodType =
           substituteTypeArgsInGenericMethodType(
@@ -830,7 +885,7 @@ public static Nullness getGenericMethodReturnTypeNullness(
    * @return Nullness of invocation's return type, or {@code NONNULL} if the call does not invoke an
    *     instance method
    */
-  public static Nullness getGenericReturnNullnessAtInvocation(
+  public Nullness getGenericReturnNullnessAtInvocation(
       Symbol.MethodSymbol invokedMethodSymbol,
       MethodInvocationTree tree,
       VisitorState state,
@@ -883,7 +938,7 @@ private static com.sun.tools.javac.util.List<Type> convertTreesToTypes(
    * @param config the NullAway config
    * @return the substituted method type for the generic method
    */
-  private static Type substituteTypeArgsInGenericMethodType(
+  private Type substituteTypeArgsInGenericMethodType(
       MethodInvocationTree methodInvocationTree,
       Symbol.MethodSymbol methodSymbol,
       VisitorState state,
@@ -894,6 +949,17 @@ private static Type substituteTypeArgsInGenericMethodType(
 
     Type.ForAll forAllType = (Type.ForAll) methodSymbol.type;
     Type.MethodType underlyingMethodType = (Type.MethodType) forAllType.qtype;
+
+    // There are no explicit type arguments, so use the inferred types
+    if (explicitTypeArgs.isEmpty()) {
+      if (inferredSubstitutionsForGenericMethodCalls.containsKey(methodInvocationTree)) {
+        return substituteInferredTypesForTypeVariables(
+            state,
+            underlyingMethodType,
+            inferredSubstitutionsForGenericMethodCalls.get(methodInvocationTree),
+            config);
+      }
+    }
     return TypeSubstitutionUtils.subst(
         state.getTypes(), underlyingMethodType, forAllType.tvars, explicitTypeArgs, config);
   }
@@ -932,7 +998,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,
@@ -941,7 +1007,6 @@ public static Nullness getGenericParameterNullnessAtInvocation(
     // If generic method invocation
     if (!invokedMethodSymbol.getTypeParameters().isEmpty()) {
       // Substitute the argument types within the MethodType
-      // NOTE: if explicitTypeArgs is empty, this is a noop
       List<Type> substitutedParamTypes =
           substituteTypeArgsInGenericMethodType(tree, invokedMethodSymbol, state, config)
               .getParameterTypes();
@@ -1160,6 +1225,14 @@ public static boolean passingLambdaOrMethodRefWithGenericReturnToUnmarkedCode(
     return callingUnannotated;
   }
 
+  /**
+   * Clears the cache of inferred substitutions for generic method calls. This should be invoked
+   * after each CompilationUnit to avoid memory leaks.
+   */
+  public void clearCache() {
+    inferredSubstitutionsForGenericMethodCalls.clear();
+  }
+
   public static boolean isNullableAnnotated(Type type, Config config) {
     return Nullness.hasNullableAnnotation(type.getAnnotationMirrors().stream(), config);
   }
diff --git a/nullaway/src/main/java/com/uber/nullaway/generics/InferSubstitutionViaAssignmentContextVisitor.java b/nullaway/src/main/java/com/uber/nullaway/generics/InferSubstitutionViaAssignmentContextVisitor.java
diff --git a/nullaway/src/test/java/com/uber/nullaway/jspecify/GenericMethodTests.java b/nullaway/src/test/java/com/uber/nullaway/jspecify/GenericMethodTests.java
