JSpecify: improve inference for generic methods based on method reference arguments

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

@@ -124,6 +124,7 @@
 import org.checkerframework.nullaway.dataflow.cfg.node.MethodInvocationNode;
 import org.checkerframework.nullaway.javacutil.ElementUtils;
 import org.checkerframework.nullaway.javacutil.TreeUtils;
+import org.jetbrains.annotations.Contract;
 import org.jspecify.annotations.Nullable;
 
 /**
@@ -808,6 +809,12 @@ private Description checkParamOverriding(
         (overridingMethod != null
                 && !codeAnnotationInfo.isSymbolUnannotated(overridingMethod, config, handler))
             || lambdaExpressionTree != null;
+    Type.MethodType jspecifyMemberReferenceMethodType = null;
+    if (memberReferenceTree != null) {
+      jspecifyMemberReferenceMethodType =
+          genericsChecks.getMemberReferenceMethodType(
+              memberReferenceTree, castToNonNull(overridingMethod), state);
+    }
 
     // Get argument nullability for the overridden method.  If overriddenMethodArgNullnessMap[i] is
     // null, parameter i is treated as unannotated.
@@ -897,6 +904,14 @@ private Description checkParamOverriding(
       }
       int methodParamInd = i - startParam;
       VarSymbol paramSymbol = overridingParamSymbols.get(methodParamInd);
+      boolean paramIsNonNull =
+          paramOfOverridingMethodIsNonNull(
+              paramSymbol,
+              methodParamInd,
+              overridingMethod,
+              isOverridingMethodAnnotated,
+              memberReferenceTree,
+              jspecifyMemberReferenceMethodType);
       // in the case where we have a parameter of a lambda expression, we do
       // *not* force the parameter to be annotated with @Nullable; instead we "inherit"
       // nullability from the corresponding functional interface method.
@@ -906,7 +921,7 @@ private Description checkParamOverriding(
           lambdaExpressionTree != null
               && NullabilityUtil.lambdaParamIsImplicitlyTyped(
                   lambdaExpressionTree.getParameters().get(methodParamInd));
-      if (!implicitlyTypedLambdaParam && paramIsNonNull(paramSymbol, isOverridingMethodAnnotated)) {
+      if (!implicitlyTypedLambdaParam && paramIsNonNull) {
         String message =
             "parameter "
                 + paramSymbol.name.toString()
@@ -945,14 +960,94 @@ private Description checkParamOverriding(
     return Description.NO_MATCH;
   }
 
-  private boolean paramIsNonNull(VarSymbol paramSymbol, boolean isMethodAnnotated) {
+  /**
+   * Checks if the parameter of the overriding method is {@code @NonNull}
+   *
+   * @param paramSymbol the symbol for the parameter of the overriding method
+   * @param methodParamInd the index of the parameter in the method signature of the overriding
+   *     method (adjusted for unbound member references)
+   * @param overridingMethod if available, the symbol for the overriding method
+   * @param isMethodAnnotated whether the overriding method is annotated
+   * @param memberReferenceTree if the overriding method is a member reference, the tree for the
+   *     member reference; otherwise {@code null}
+   * @param memberReferenceMethodType if the overriding method is a member reference, the method
+   *     type of the member reference after handling generics; otherwise {@code null}
+   * @return true if the parameter of the overriding method is effectively {@code @NonNull}, false
+   *     otherwise
+   */
+  private boolean paramOfOverridingMethodIsNonNull(
+      VarSymbol paramSymbol,
+      int methodParamInd,
+      Symbol.@Nullable MethodSymbol overridingMethod,
+      boolean isMethodAnnotated,
+      @Nullable MemberReferenceTree memberReferenceTree,
+      Type.@Nullable MethodType memberReferenceMethodType) {
+    boolean result = false;
     if (isMethodAnnotated) {
-      return !Nullness.hasNullableAnnotation(paramSymbol, config);
+      result = !Nullness.hasNullableAnnotation(paramSymbol, config);
     } else if (config.acknowledgeRestrictiveAnnotations()) {
       // can still be @NonNull if there is a restrictive annotation
-      return Nullness.hasNonNullAnnotation(paramSymbol, config);
+      result = Nullness.hasNonNullAnnotation(paramSymbol, config);
+    }
+    if (result && memberReferenceMethodType != null) {
+      // when the overriding method is a member reference, also check that the parameter is not
+      // effectively @Nullable due to generics.  memberReferenceMethodType should be the method type
+      // of the member reference after handling generics
+      com.sun.tools.javac.util.List<Type> parameterTypes =
+          memberReferenceMethodType.getParameterTypes();
+      int memberRefParamIndex =
+          getMemberRefParamIndex(methodParamInd, overridingMethod, memberReferenceTree);
+      Type paramType = parameterTypes.get(memberRefParamIndex);
+      // if we have a method reference to a varargs method where varargs are passed individually
+      // (not as an array), and this parameter is the varargs
+      // parameter, then we need to check the nullability of the varargs element type
+      if (memberRefToVarargsPassedIndividually(overridingMethod, memberReferenceTree)
+          && methodParamInd >= overridingMethod.getParameters().size() - 1) {
+        Verify.verify(
+            paramType.getKind() == TypeKind.ARRAY,
+            "Expected array type for varargs parameter in %s, got %s",
+            memberReferenceTree,
+            paramType);
+        paramType = ((Type.ArrayType) paramType).getComponentType();
+      }
+      result = !Nullness.hasNullableAnnotation(paramType.getAnnotationMirrors().stream(), config);
     }
-    return false;
+    return result;
+  }
+
+  /**
+   * Checks if we have a member reference to a varargs method where the varargs are passed
+   * individually rather than as an array
+   *
+   * @param referencedMethod if available, the symbol for the referenced method
+   * @param memberReferenceTree the tree for the member reference
+   * @return true if we have a member reference to a varargs method where the varargs are passed
+   *     individually
+   */
+  @Contract("null, _ -> false; _, null -> false")
+  private static boolean memberRefToVarargsPassedIndividually(
+      Symbol.@Nullable MethodSymbol referencedMethod,
+      @Nullable MemberReferenceTree memberReferenceTree) {
+    return memberReferenceTree != null
+        && referencedMethod != null
+        && referencedMethod.isVarArgs()
+        && ((JCTree.JCMemberReference) memberReferenceTree).varargsElement != null;
+  }
+
+  private static int getMemberRefParamIndex(
+      int methodParamInd,
+      Symbol.@Nullable MethodSymbol overridingMethod,
+      @Nullable MemberReferenceTree memberReferenceTree) {
+    int memberRefParamIndex = methodParamInd;
+    if (memberRefToVarargsPassedIndividually(overridingMethod, memberReferenceTree)) {
+      // With varargs adaptation, one or more functional-interface parameters can map to the
+      // varargs element type of the referenced method.
+      int varargsParamIndex = overridingMethod.getParameters().size() - 1;
+      if (methodParamInd >= varargsParamIndex) {
+        memberRefParamIndex = varargsParamIndex;
+      }
+    }
+    return memberRefParamIndex;
   }
 
   static Trees getTreesInstance(VisitorState state) {

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

@@ -852,10 +852,6 @@ private void generateConstraintsForPseudoAssignment(
       ExpressionTree rhsExpr,
       Type lhsType) {
     rhsExpr = ASTHelpers.stripParentheses(rhsExpr);
-    if (rhsExpr instanceof MemberReferenceTree) {
-      // TODO generate constraints from method reference argument types
-      return;
-    }
     // if the parameter is itself a generic call requiring inference, generate constraints for
     // that call
     if (isGenericCallNeedingInference(rhsExpr)) {
@@ -864,16 +860,18 @@ private void generateConstraintsForPseudoAssignment(
       allInvocations.add(invTree);
       generateConstraintsForCall(
           state, path, lhsType, false, solver, symbol, invTree, allInvocations);
-    } else if (!(rhsExpr instanceof LambdaExpressionTree lambda)) {
+    } else if (rhsExpr instanceof LambdaExpressionTree lambda) {
+      handleLambdaInGenericMethodInference(state, path, solver, allInvocations, lhsType, lambda);
+    } else if (rhsExpr instanceof MemberReferenceTree memberReferenceTree) {
+      handleMethodRefInGenericMethodInference(state, solver, lhsType, memberReferenceTree);
+    } else { // all other cases
       Type argumentType = getTreeType(rhsExpr, state);
       if (argumentType == null) {
         // bail out of any checking involving raw types for now
         return;
       }
       argumentType = refineArgumentTypeWithDataflow(argumentType, rhsExpr, state, path);
       solver.addSubtypeConstraint(argumentType, lhsType, false);
-    } else {
-      handleLambdaInGenericMethodInference(state, path, solver, allInvocations, lhsType, lambda);
     }
   }
 
@@ -928,6 +926,175 @@ private void handleLambdaInGenericMethodInference(
     }
   }
 
+  /**
+   * Generate constraints for a method reference argument by comparing functional interface method
+   * parameter and return types against the referenced method.
+   *
+   * @param state the visitor state
+   * @param solver the constraint solver
+   * @param lhsType the type to which the method reference is being assigned
+   * @param memberReferenceTree the method reference argument
+   */
+  private void handleMethodRefInGenericMethodInference(
+      VisitorState state,
+      ConstraintSolver solver,
+      Type lhsType,
+      MemberReferenceTree memberReferenceTree) {
+    if (lhsType.isRaw()) {
+      return;
+    }
+    Types types = state.getTypes();
+
+    // first, figure out the proper method type to use for the member reference
+    Symbol.MethodSymbol referencedMethod = ASTHelpers.getSymbol(memberReferenceTree);
+    if (referencedMethod == null || referencedMethod.isConstructor()) {
+      // TODO handle constructor references like Foo::new;
+      //  https://github.com/uber/NullAway/issues/1468
+      return;
+    }
+    Type.MethodType referencedMethodType = referencedMethod.asType().asMethodType();
+    Type qualifierType = null;
+    if (!referencedMethod.isStatic()) {
+      // get the referenced method type as a member of the type of the qualifier expression
+      qualifierType = getTreeType(memberReferenceTree.getQualifierExpression(), state);
+      if (qualifierType != null) {
+        referencedMethodType =
+            TypeSubstitutionUtils.memberType(types, qualifierType, referencedMethod, config)
+                .asMethodType();
+      }
+    }
+    // substitute any explicit type arguments from the member reference
+    List<? extends ExpressionTree> typeArgumentTrees = memberReferenceTree.getTypeArguments();
+    if (typeArgumentTrees != null && !typeArgumentTrees.isEmpty()) {
+      referencedMethodType =
+          TypeSubstitutionUtils.subst(
+                  state.getTypes(),
+                  referencedMethodType,
+                  ((Type.ForAll) referencedMethod.asType()).tvars,
+                  convertTreesToTypes(typeArgumentTrees),
+                  config)
+              .asMethodType();
+    }
+    // allow for handler overrides
+    referencedMethodType =
+        handler.onOverrideMethodType(referencedMethod, referencedMethodType, state);
+
+    // now, get the type of the corresponding functional interface method, as a member of lhsType
+    Symbol.MethodSymbol fiMethod =
+        NullabilityUtil.getFunctionalInterfaceMethod(memberReferenceTree, types);
+    Type.MethodType fiMethodTypeAsMember =
+        TypeSubstitutionUtils.memberType(types, lhsType, fiMethod, config).asMethodType();
+
+    // constrain returns: method reference return type <: functional interface return type
+    Type fiReturnType = fiMethodTypeAsMember.getReturnType();
+    Type referencedReturnType = referencedMethodType.getReturnType();
+    if (fiReturnType.getKind() != TypeKind.VOID
+        && referencedReturnType.getKind() != TypeKind.VOID) {
+      solver.addSubtypeConstraint(referencedReturnType, fiReturnType, false);
+    }
+
+    // constrain parameters:
+    //  i^{th} functional interface parameter type <: i^{th} method reference parameter type,
+    //  aligned appropriately in the case of unbound method references
+    com.sun.tools.javac.util.List<Type> fiParamTypes = fiMethodTypeAsMember.getParameterTypes();
+    com.sun.tools.javac.util.List<Type> referencedParamTypes =
+        referencedMethodType.getParameterTypes();
+    int fiStartIndex = 0;
+    if (((JCTree.JCMemberReference) memberReferenceTree).kind.isUnbound()) {
+      // an unbound method reference like String::length has an implicit receiver parameter
+      // that needs to be aligned with the first functional interface parameter
+      Verify.verify(
+          !fiParamTypes.isEmpty(),
+          "Expected receiver parameter for unbound method ref %s",
+          memberReferenceTree);
+      if (qualifierType != null) {
+        Type receiverType = fiParamTypes.get(0);
+        solver.addSubtypeConstraint(receiverType, qualifierType, false);
+      }
+      fiStartIndex = 1;
+    }
+    // first, handle the non-varargs case
+    int fiParamCount = fiParamTypes.size() - fiStartIndex;
+    int nonVarargsParamCount =
+        referencedMethod.isVarArgs()
+            ? Math.min(fiParamCount, referencedParamTypes.size() - 1)
+            : referencedParamTypes.size();
+    for (int i = 0; i < nonVarargsParamCount; i++) {
+      solver.addSubtypeConstraint(
+          fiParamTypes.get(fiStartIndex + i), referencedParamTypes.get(i), false);
+    }
+    if (!referencedMethod.isVarArgs()) {
+      return;
+    }
+
+    // For varargs references, the functional interface can map to fixed-arity form (single array
+    // argument at the varargs position) or variable-arity form (zero or more element arguments).
+    int varargsParamPosition = referencedParamTypes.size() - 1;
+    if (fiParamCount == varargsParamPosition) {
+      // No varargs arguments; this is the variable-arity case, passing zero arguments
+      return;
+    }
+    Type varargsArrayType = referencedParamTypes.get(varargsParamPosition);
+    Verify.verify(
+        varargsArrayType.getKind() == TypeKind.ARRAY,
+        "Expected array type for varargs parameter in %s, got %s",
+        memberReferenceTree,
+        varargsArrayType);
+    JCTree.JCMemberReference javacMemberRef = (JCTree.JCMemberReference) memberReferenceTree;
+    int firstVarargsFiParamIndex = fiStartIndex + varargsParamPosition;
+    if (javacMemberRef.varargsElement == null) {
+      // javac resolved this member reference using non-varargs (fixed-arity) adaptation.
+      solver.addSubtypeConstraint(
+          fiParamTypes.get(firstVarargsFiParamIndex), varargsArrayType, false);
+      return;
+    }
+    // javac resolved this member reference using varargs (variable-arity) adaptation.
+    // Use the element type from the referenced varargs array type
+    Type varargsElementType = types.elemtype(varargsArrayType);
+    for (int i = varargsParamPosition; i < fiParamCount; i++) {
+      solver.addSubtypeConstraint(fiParamTypes.get(fiStartIndex + i), varargsElementType, false);
+    }
+  }
+
+  /**
+   * Gets the method type for a member reference handling generics, in JSpecify mode
+   *
+   * @param memberReferenceTree the member reference tree
+   * @param overridingMethod the method symbol for the method referenced by {@code
+   *     memberReferenceTree}
+   * @param state the visitor state
+   * @return the method type for the member reference, with generics handled, or null if not in
+   *     JSpecify mode
+   */
+  public Type.@Nullable MethodType getMemberReferenceMethodType(
+      MemberReferenceTree memberReferenceTree,
+      Symbol.MethodSymbol overridingMethod,
+      VisitorState state) {
+    if (!config.isJSpecifyMode()) {
+      return null;
+    }
+    Type.MethodType result = overridingMethod.asType().asMethodType();
+    if (!overridingMethod.isStatic()) {
+      // This handles any generic type parameters of the qualifier of the member reference, e.g. for
+      // x::m, where x is of type Foo<Integer>, it handles the type parameter Integer whereever it
+      // appears in the signature of m.
+      Type qualifierType = ASTHelpers.getType(memberReferenceTree.getQualifierExpression());
+      if (qualifierType != null && !qualifierType.isRaw()) {
+        result =
+            TypeSubstitutionUtils.memberType(
+                    state.getTypes(), qualifierType, overridingMethod, config)
+                .asMethodType();
+      }
+    }
+    if (overridingMethod.asType() instanceof Type.ForAll) {
+      // the referenced method is a generic method; we need to substitute inferred nullability for
+      // type arguments if it was inferred
+      result = getInferredMethodTypeForGenericMethodReference(result, state);
+    }
+    // finally, run any handlers
+    return handler.onOverrideMethodType(overridingMethod, result, state);
+  }
+
   /**
    * A visitor that scans a {@link Tree} (typically a lambda or method body) to find all {@code
    * return} statements and collect their expressions.
@@ -1365,6 +1532,36 @@ public void compareGenericTypeParameterNullabilityForCall(
             });
   }
 
+  /**
+   * For a generic method reference, if it is being called in a context that requires type argument
+   * nullability inference, return the method type with inferred nullability for type parameters.
+   * Otherwise, return the original method type.
+   *
+   * @param methodType the original method type
+   * @param state the visitor state (generic method reference should be leaf of {@code
+   *     state.getPath()})
+   * @return the method type with inferred nullability for type parameters if inference was
+   *     performed, or the original method type otherwise
+   */
+  private Type.MethodType getInferredMethodTypeForGenericMethodReference(
+      Type.MethodType methodType, VisitorState state) {
+    TreePath parentPath = state.getPath().getParentPath();
+    while (parentPath != null && parentPath.getLeaf() instanceof ParenthesizedTree) {
+      parentPath = parentPath.getParentPath();
+    }
+    Tree parentTree = parentPath != null ? parentPath.getLeaf() : null;
+    if (parentTree instanceof MethodInvocationTree methodInvocationTree
+        && isGenericCallNeedingInference(methodInvocationTree)) {
+      MethodInferenceResult inferenceResult =
+          inferredTypeVarNullabilityForGenericCalls.get(methodInvocationTree);
+      if (inferenceResult instanceof InferenceSuccess successResult) {
+        return TypeSubstitutionUtils.updateMethodTypeWithInferredNullability(
+            methodType, methodType, successResult.typeVarNullability, state, config);
+      }
+    }
+    return methodType;
+  }
+
   /**
    * Checks that type parameter nullability is consistent between an overriding method and the
    * corresponding overridden method.