Inference of generic method type arguments based on returns and parameter passing (#1226)

msridhar · web-flow · commit 20dd1526c96d · 2025-06-30T06:50:31.000+05:30
diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -985,7 +985,7 @@ private Description checkReturnExpression(
 
     // Check generic type arguments for returned expression here, since we need to check the type
     // arguments regardless of the top-level nullability of the return type
-    GenericsChecks.checkTypeParameterNullnessForFunctionReturnType(
+    genericsChecks.checkTypeParameterNullnessForFunctionReturnType(
         retExpr, methodSymbol, this, state);
 
     // Now, perform the check for returning @Nullable from @NonNull.  First, we check if the return
@@ -1962,7 +1962,7 @@ private Description handleInvocation(
         }
         actual = actualParams.get(argPos);
         VarSymbol formalParamSymbol = formalParams.get(formalParams.size() - 1);
-        boolean isVarArgsCall = isVarArgsCall(tree);
+        boolean isVarArgsCall = NullabilityUtil.isVarArgsCall(tree);
         if (isVarArgsCall) {
           // This is the case were varargs are being passed individually, as 1 or more actual
           // arguments starting at the position of the var args formal.
@@ -2016,23 +2016,6 @@ private Description handleInvocation(
     return checkCastToNonNullTakesNullable(tree, state, methodSymbol, actualParams);
   }
 
-  /**
-   * Checks if the method invocation is a varargs call, i.e., if individual arguments are being
-   * passed in the varargs position. If false, it means that an array is being passed in the varargs
-   * position.
-   *
-   * @param tree the method invocation tree (MethodInvocationTree or NewClassTree)
-   * @return true if the method invocation is a varargs call, false otherwise
-   */
-  private boolean isVarArgsCall(Tree tree) {
-    // javac sets the varargsElement field to a non-null value if the invocation is a varargs call
-    Type varargsElement =
-        tree instanceof JCTree.JCMethodInvocation
-            ? ((JCTree.JCMethodInvocation) tree).varargsElement
-            : ((JCTree.JCNewClass) tree).varargsElement;
-    return varargsElement != null;
-  }
-
   private Description checkCastToNonNullTakesNullable(
       Tree tree,
       VisitorState state,
diff --git a/nullaway/src/main/java/com/uber/nullaway/NullabilityUtil.java b/nullaway/src/main/java/com/uber/nullaway/NullabilityUtil.java
@@ -645,4 +645,21 @@ public static boolean hasJetBrainsNotNullDeclarationAnnotation(Symbol.VarSymbol
         .map(a -> a.getAnnotationType().toString())
         .anyMatch(annotName -> annotName.equals(JETBRAINS_NOT_NULL));
   }
+
+  /**
+   * Checks if the method invocation is a varargs call, i.e., if individual arguments are being
+   * passed in the varargs position. If false, it means that an array is being passed in the varargs
+   * position.
+   *
+   * @param tree the method invocation tree (MethodInvocationTree or NewClassTree)
+   * @return true if the method invocation is a varargs call, false otherwise
+   */
+  public static boolean isVarArgsCall(Tree tree) {
+    // javac sets the varargsElement field to a non-null value if the invocation is a varargs call
+    Type varargsElement =
+        tree instanceof JCTree.JCMethodInvocation
+            ? ((JCTree.JCMethodInvocation) tree).varargsElement
+            : ((JCTree.JCNewClass) tree).varargsElement;
+    return varargsElement != null;
+  }
 }
diff --git a/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java b/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java
@@ -33,6 +33,7 @@
 import com.uber.nullaway.ErrorBuilder;
 import com.uber.nullaway.ErrorMessage;
 import com.uber.nullaway.NullAway;
+import com.uber.nullaway.NullabilityUtil;
 import com.uber.nullaway.Nullness;
 import com.uber.nullaway.handlers.Handler;
 import java.util.ArrayList;
@@ -446,42 +447,92 @@ public void checkTypeParameterNullnessForAssignability(
       return;
     }
     Type rhsType = getTreeType(rhsTree, config);
-
-    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
-        boolean invokedMethodIsNullUnmarked =
-            CodeAnnotationInfo.instance(state.context)
-                .isSymbolUnannotated(methodSymbol, config, analysis.getHandler());
-        InferGenericMethodSubstitutionViaAssignmentContextVisitor inferVisitor =
-            new InferGenericMethodSubstitutionViaAssignmentContextVisitor(
-                state, config, invokedMethodIsNullUnmarked);
-        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) {
+      if (rhsTree instanceof MethodInvocationTree) {
+        rhsType =
+            inferGenericMethodCallType(
+                analysis, state, (MethodInvocationTree) rhsTree, config, lhsType, rhsType);
+      }
       boolean isAssignmentValid = subtypeParameterNullability(lhsType, rhsType, state, config);
       if (!isAssignmentValid) {
         reportInvalidAssignmentInstantiationError(tree, lhsType, rhsType, state, analysis);
       }
     }
   }
 
+  /**
+   * Infers the type of a generic method call based on the assignment context. Side-effects the
+   * #inferredSubstitutionsForGenericMethodCalls map with the inferred type.
+   *
+   * @param analysis the analysis
+   * @param state the visitor state
+   * @param invocationTree the method invocation tree representing the call to a generic method
+   * @param config the analysis config
+   * @param typeFromAssignmentContext the type being "assigned to" in the assignment context
+   * @param exprType the type of the right-hand side of the pseudo-assignment, which may be null
+   * @return the type of the method call after inference
+   */
+  private Type inferGenericMethodCallType(
+      NullAway analysis,
+      VisitorState state,
+      MethodInvocationTree invocationTree,
+      Config config,
+      Type typeFromAssignmentContext,
+      Type exprType) {
+    Type result = exprType;
+    MethodInvocationTree methodInvocationTree = invocationTree;
+    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
+      boolean invokedMethodIsNullUnmarked =
+          CodeAnnotationInfo.instance(state.context)
+              .isSymbolUnannotated(methodSymbol, config, analysis.getHandler());
+      Map<TypeVariable, Type> substitution;
+      Type returnType = methodSymbol.getReturnType();
+      if (returnType instanceof Type.TypeVar) {
+        // we need different logic if the return type is a type variable
+        // if the assignment context type is @Nullable, we shouldn't infer anything, since that
+        // accommodates the type argument being either @Nullable or @NonNull
+        Type.TypeVar typeVar = (Type.TypeVar) returnType;
+        substitution = new LinkedHashMap<>();
+        boolean nonNullAssignmentContextType =
+            !Nullness.hasNullableAnnotation(
+                typeFromAssignmentContext.getAnnotationMirrors().stream(), config);
+        if (nonNullAssignmentContextType) {
+          // if the assignment context type is @NonNull, we can just use it
+          substitution.put(typeVar, typeFromAssignmentContext);
+        } else {
+          Type upperBound = typeVar.getUpperBound();
+          boolean typeVarHasNullableUpperBound =
+              Nullness.hasNullableAnnotation(upperBound.getAnnotationMirrors().stream(), config);
+          // if the type variable cannot be @Nullable, we can use the lhsType with any @Nullable
+          // annotation stripped
+          if (!typeVarHasNullableUpperBound && !invokedMethodIsNullUnmarked) {
+            // we can use the lhsType with any @Nullable annotation stripped
+            // TODO we should just strip out the top-level @Nullable annotation;
+            //  stripMetadata() also removes nested @Nullable annotations
+            substitution.put(typeVar, typeFromAssignmentContext.stripMetadata());
+          }
+        }
+
+      } else {
+        InferGenericMethodSubstitutionViaAssignmentContextVisitor inferVisitor =
+            new InferGenericMethodSubstitutionViaAssignmentContextVisitor(
+                state, config, invokedMethodIsNullUnmarked);
+        returnType.accept(inferVisitor, typeFromAssignmentContext);
+        substitution = inferVisitor.getInferredSubstitution();
+      }
+      inferredSubstitutionsForGenericMethodCalls.put(methodInvocationTree, substitution);
+      // update with inferred substitution
+      result =
+          substituteInferredTypesForTypeVariables(
+              state, methodSymbol.getReturnType(), substitution, config);
+    }
+    return result;
+  }
+
   /**
    * Substitutes inferred types for type variables within a type.
    *
@@ -512,7 +563,7 @@ private Type substituteInferredTypesForTypeVariables(
    * @param analysis the analysis object
    * @param state the visitor state
    */
-  public static void checkTypeParameterNullnessForFunctionReturnType(
+  public void checkTypeParameterNullnessForFunctionReturnType(
       ExpressionTree retExpr,
       Symbol.MethodSymbol methodSymbol,
       NullAway analysis,
@@ -529,6 +580,16 @@ public static void checkTypeParameterNullnessForFunctionReturnType(
     }
     Type returnExpressionType = getTreeType(retExpr, config);
     if (formalReturnType != null && returnExpressionType != null) {
+      if (retExpr instanceof MethodInvocationTree) {
+        returnExpressionType =
+            inferGenericMethodCallType(
+                analysis,
+                state,
+                (MethodInvocationTree) retExpr,
+                config,
+                formalReturnType,
+                returnExpressionType);
+      }
       boolean isReturnTypeValid =
           subtypeParameterNullability(formalReturnType, returnExpressionType, state, config);
       if (!isReturnTypeValid) {
@@ -730,28 +791,48 @@ public void compareGenericTypeParameterNullabilityForCall(
         // bail out of any checking involving raw types for now
         return;
       }
-      Type actualParameter = getTreeType(actualParams.get(i), config);
-      if (actualParameter != null) {
-        if (!subtypeParameterNullability(formalParameter, actualParameter, state, config)) {
+      ExpressionTree currentActualParam = actualParams.get(i);
+      Type actualParameterType = getTreeType(currentActualParam, config);
+      if (actualParameterType != null) {
+        if (currentActualParam instanceof MethodInvocationTree) {
+          // infer the type of the method call based on the assignment context
+          // and the formal parameter type
+          actualParameterType =
+              inferGenericMethodCallType(
+                  analysis,
+                  state,
+                  (MethodInvocationTree) currentActualParam,
+                  config,
+                  formalParameter,
+                  actualParameterType);
+        }
+        if (!subtypeParameterNullability(formalParameter, actualParameterType, state, config)) {
           reportInvalidParametersNullabilityError(
-              formalParameter, actualParameter, actualParams.get(i), state, analysis);
+              formalParameter, actualParameterType, currentActualParam, state, analysis);
         }
       }
     }
-    if (isVarArgs && !formalParamTypes.isEmpty()) {
-      Type.ArrayType varargsArrayType =
-          (Type.ArrayType) formalParamTypes.get(formalParamTypes.size() - 1);
-      Type varargsElementType = varargsArrayType.elemtype;
+    if (isVarArgs && !formalParamTypes.isEmpty() && NullabilityUtil.isVarArgsCall(tree)) {
+      Type varargsElementType =
+          ((Type.ArrayType) formalParamTypes.get(formalParamTypes.size() - 1)).elemtype;
       for (int i = formalParamTypes.size() - 1; i < actualParams.size(); i++) {
-        Type actualParameterType = getTreeType(actualParams.get(i), config);
-        // If the actual parameter type is assignable to the varargs array type, then the call site
-        // is passing the varargs directly in an array, and we should skip our check.
-        if (actualParameterType != null
-            && !state.getTypes().isAssignable(actualParameterType, varargsArrayType)) {
+        ExpressionTree actualParamExpr = actualParams.get(i);
+        Type actualParameterType = getTreeType(actualParamExpr, config);
+        if (actualParameterType != null) {
+          if (actualParamExpr instanceof MethodInvocationTree) {
+            actualParameterType =
+                inferGenericMethodCallType(
+                    analysis,
+                    state,
+                    (MethodInvocationTree) actualParamExpr,
+                    config,
+                    varargsElementType,
+                    actualParameterType);
+          }
           if (!subtypeParameterNullability(
               varargsElementType, actualParameterType, state, config)) {
             reportInvalidParametersNullabilityError(
-                varargsElementType, actualParameterType, actualParams.get(i), state, analysis);
+                varargsElementType, actualParameterType, actualParamExpr, state, analysis);
           }
         }
       }
diff --git a/nullaway/src/test/java/com/uber/nullaway/jspecify/GenericMethodTests.java b/nullaway/src/test/java/com/uber/nullaway/jspecify/GenericMethodTests.java
