Support marking method type variable upper bounds as @Nullable in library models (#1345)

We need this support to address #1157 (which I'll do in a follow-up PR).
Now `LibraryModels` can mark the upper bound of some method type
variable as `@Nullable`. Rename method in `LibraryModels` for getting
upper bounds of class type variables for clarity.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Enhanced null-checking capabilities to support nullable upper bounds
in generic type variables for library methods.

* **Tests**
* Added test coverage for validating nullable upper bounds in generic
method type parameters.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: msridhar

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

@@ -115,15 +115,26 @@ public interface LibraryModels {
 
   /**
    * Get the (className, type argument index) pairs for library classes where the generic type
-   * argument has a {@code @Nullable} upper bound. Only used in JSpecify mode.
+   * variable has a {@code @Nullable} upper bound. Only used in JSpecify mode.
    *
-   * @return map from the className to the positions of the generic type arguments that have a
+   * @return map from the className to the positions of the generic type variables that have a
    *     {@code Nullable} upper bound.
    */
   default ImmutableSetMultimap<String, Integer> typeVariablesWithNullableUpperBounds() {
     return ImmutableSetMultimap.of();
   }
 
+  /**
+   * Get the (method, type argument index) pairs for library methods where the generic type variable
+   * has a {@code @Nullable} upper bound. Only used in JSpecify mode
+   *
+   * @return map from the method to the positions of the generic type variables that have a {@code
+   *     Nullable} upper bound.
+   */
+  default ImmutableSetMultimap<MethodRef, Integer> methodTypeVariablesWithNullableUpperBounds() {
+    return ImmutableSetMultimap.of();
+  }
+
   /**
    * Get the set of library classes that are NullMarked. Only used in JSpecify mode.
    *

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

@@ -299,6 +299,29 @@ private boolean upperBoundIsNullable(Element typeVarElement) {
     if (fromUnannotatedMethod(typeVarElement)) {
       return true;
     }
+    // first, check if library model overrides the upper bound nullability
+    Element enclosingElement = typeVarElement.getEnclosingElement();
+    if (enclosingElement instanceof Symbol.MethodSymbol) {
+      Symbol.MethodSymbol methodSymbol = (Symbol.MethodSymbol) enclosingElement;
+      int typeVarIndex =
+          methodSymbol.getTypeParameters().indexOf((Symbol.TypeVariableSymbol) typeVarElement);
+      // TODO typeVarIndex is -1 in some cases; see test
+      //  com.uber.nullaway.jspecify.GenericMethodTests.instanceGenericMethodWithMethodRefArgument.
+      //  Investigate further.
+      if (typeVarIndex >= 0
+          && handler.onOverrideMethodTypeVariableUpperBound(methodSymbol, typeVarIndex)) {
+        return true;
+      }
+    } else if (enclosingElement instanceof Symbol.ClassSymbol) {
+      Symbol.ClassSymbol classSymbol = (Symbol.ClassSymbol) enclosingElement;
+      int typeVarIndex =
+          classSymbol.getTypeParameters().indexOf((Symbol.TypeVariableSymbol) typeVarElement);
+      if (typeVarIndex >= 0
+          && handler.onOverrideClassTypeVariableUpperBound(classSymbol.toString(), typeVarIndex)) {
+        return true;
+      }
+    }
+    // otherwise, check the actual upper bound annotations
     Type upperBound = (Type) ((TypeVariable) typeVarElement.asType()).getUpperBound();
     com.sun.tools.javac.util.List<Attribute.TypeCompound> annotationMirrors =
         upperBound.getAnnotationMirrors();

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

@@ -179,7 +179,7 @@ private boolean[] getTypeParamsWithNullableUpperBound(Type type) {
       com.sun.tools.javac.util.List<Attribute.TypeCompound> annotationMirrors =
           upperBound.getAnnotationMirrors();
       if (Nullness.hasNullableAnnotation(annotationMirrors.stream(), config)
-          || handler.onOverrideTypeParameterUpperBound(type.tsym.toString(), i)) {
+          || handler.onOverrideClassTypeVariableUpperBound(type.tsym.toString(), i)) {
         result[i] = true;
       }
     }
@@ -254,10 +254,11 @@ public void checkGenericMethodCallTypeArguments(Tree tree, VisitorState state) {
             upperBound.getAnnotationMirrors();
         boolean hasNullableAnnotation =
             Nullness.hasNullableAnnotation(annotationMirrors.stream(), config)
-                || handler.onOverrideTypeParameterUpperBound(baseType.tsym.toString(), i);
+                || handler.onOverrideClassTypeVariableUpperBound(baseType.tsym.toString(), i);
         // if type variable's upper bound does not have @Nullable annotation then the instantiation
         // is invalid
-        if (!hasNullableAnnotation) {
+        if (!hasNullableAnnotation
+            && !handler.onOverrideMethodTypeVariableUpperBound(methodSymbol, i)) {
           reportInvalidTypeArgumentError(
               nullableTypeArguments.get(i), methodSymbol, typeVariable, state);
         }

nullaway/src/main/java/com/uber/nullaway/handlers/BaseNoOpHandler.java

@@ -208,7 +208,7 @@ public void onNonNullFieldAssignment(
   }
 
   @Override
-  public boolean onOverrideTypeParameterUpperBound(String className, int index) {
+  public boolean onOverrideClassTypeVariableUpperBound(String className, int index) {
     return false;
   }
 

nullaway/src/main/java/com/uber/nullaway/handlers/CompositeHandler.java

@@ -312,10 +312,24 @@ public MethodInvocationNode onCFGBuildPhase1AfterVisitMethodInvocation(
 
   /** Returns true if any handler returns true. */
   @Override
-  public boolean onOverrideTypeParameterUpperBound(String className, int index) {
+  public boolean onOverrideClassTypeVariableUpperBound(String className, int index) {
     boolean result = false;
     for (Handler h : handlers) {
-      result = h.onOverrideTypeParameterUpperBound(className, index);
+      result = h.onOverrideClassTypeVariableUpperBound(className, index);
+      if (result) {
+        break;
+      }
+    }
+    return result;
+  }
+
+  /** Returns true if any handler returns true. */
+  @Override
+  public boolean onOverrideMethodTypeVariableUpperBound(
+      Symbol.MethodSymbol methodSymbol, int index) {
+    boolean result = false;
+    for (Handler h : handlers) {
+      result = h.onOverrideMethodTypeVariableUpperBound(methodSymbol, index);
       if (result) {
         break;
       }

nullaway/src/main/java/com/uber/nullaway/handlers/Handler.java

@@ -392,7 +392,20 @@ MethodInvocationNode onCFGBuildPhase1AfterVisitMethodInvocation(
    * @return boolean true if the variable should be treated as having a {@code @Nullable} upper
    *     bound
    */
-  boolean onOverrideTypeParameterUpperBound(String className, int index);
+  boolean onOverrideClassTypeVariableUpperBound(String className, int index);
+
+  /**
+   * Method to override the nullability of the upper bound for a generic type variable on a method.
+   *
+   * @param methodSymbol The method symbol
+   * @param index index of the generic type variable (starting at 0)
+   * @return boolean true if the variable should be treated as having a {@code @Nullable} upper
+   *     bound
+   */
+  default boolean onOverrideMethodTypeVariableUpperBound(
+      Symbol.MethodSymbol methodSymbol, int index) {
+    return false;
+  }
 
   /**
    * Method to override the null-markedness of a class.

nullaway/src/main/java/com/uber/nullaway/handlers/LibraryModelsHandler.java

@@ -353,11 +353,21 @@ private void setUnconditionalArgumentNullness(
   }
 
   @Override
-  public boolean onOverrideTypeParameterUpperBound(String className, int index) {
+  public boolean onOverrideClassTypeVariableUpperBound(String className, int index) {
     ImmutableSet<Integer> res = libraryModels.typeVariablesWithNullableUpperBounds().get(className);
     return res.contains(index);
   }
 
+  @Override
+  public boolean onOverrideMethodTypeVariableUpperBound(
+      Symbol.MethodSymbol methodSymbol, int index) {
+    ImmutableSet<Integer> res =
+        libraryModels
+            .methodTypeVariablesWithNullableUpperBounds()
+            .get(MethodRef.fromSymbol(methodSymbol));
+    return res.contains(index);
+  }
+
   @Override
   public boolean onOverrideNullMarkedClasses(String className) {
     return libraryModels.nullMarkedClasses().contains(className);
@@ -1004,6 +1014,9 @@ private static class CombinedLibraryModels implements LibraryModels {
 
     private final ImmutableSetMultimap<String, Integer> nullableVariableTypeUpperBounds;
 
+    private final ImmutableSetMultimap<MethodRef, Integer>
+        methodTypeVariablesWithNullableUpperBounds;
+
     private final ImmutableSet<String> nullMarkedClasses;
 
     private final ImmutableSet<FieldRef> nullableFields;
@@ -1022,6 +1035,8 @@ private static class CombinedLibraryModels implements LibraryModels {
           new ImmutableSetMultimap.Builder<>();
       ImmutableSetMultimap.Builder<String, Integer> nullableVariableTypeUpperBoundsBuilder =
           new ImmutableSetMultimap.Builder<>();
+      ImmutableSetMultimap.Builder<MethodRef, Integer>
+          methodTypeVariableNullableUpperBoundsBuilder = new ImmutableSetMultimap.Builder<>();
       ImmutableSet.Builder<String> nullMarkedClassesBuilder = new ImmutableSet.Builder<>();
       ImmutableSetMultimap.Builder<MethodRef, Integer> nullImpliesTrueParametersBuilder =
           new ImmutableSetMultimap.Builder<>();
@@ -1097,6 +1112,8 @@ private static class CombinedLibraryModels implements LibraryModels {
         }
         nullableVariableTypeUpperBoundsBuilder.putAll(
             libraryModels.typeVariablesWithNullableUpperBounds());
+        methodTypeVariableNullableUpperBoundsBuilder.putAll(
+            libraryModels.methodTypeVariablesWithNullableUpperBounds());
         for (String name : libraryModels.nullMarkedClasses()) {
           nullMarkedClassesBuilder.add(name);
         }
@@ -1119,6 +1136,8 @@ private static class CombinedLibraryModels implements LibraryModels {
       customStreamNullabilitySpecs = customStreamNullabilitySpecsBuilder.build();
       nullableFields = nullableFieldsBuilder.build();
       nullableVariableTypeUpperBounds = nullableVariableTypeUpperBoundsBuilder.build();
+      methodTypeVariablesWithNullableUpperBounds =
+          methodTypeVariableNullableUpperBoundsBuilder.build();
       nullMarkedClasses = nullMarkedClassesBuilder.build();
     }
 
@@ -1171,6 +1190,11 @@ public ImmutableSetMultimap<String, Integer> typeVariablesWithNullableUpperBound
       return nullableVariableTypeUpperBounds;
     }
 
+    @Override
+    public ImmutableSetMultimap<MethodRef, Integer> methodTypeVariablesWithNullableUpperBounds() {
+      return methodTypeVariablesWithNullableUpperBounds;
+    }
+
     @Override
     public ImmutableSet<String> nullMarkedClasses() {
       return nullMarkedClasses;

nullaway/src/test/java/com/uber/nullaway/jspecify/GenericMethodTests.java

@@ -1448,6 +1448,24 @@ public void issue1294_lambdaArguments() {
         .doTest();
   }
 
+  @Test
+  public void instanceGenericMethodWithMethodRefArgument() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            "import org.jspecify.annotations.NullMarked;",
+            "import java.util.List;",
+            "import java.util.function.Consumer;",
+            "@NullMarked",
+            "class Test {",
+            "        public <E extends Enum<E>> void visitEnum(String descriptor, String value, Consumer<E> consumer) {}",
+            "        void test(String s1, String s2, List<Object> l) {",
+            "            visitEnum(s1, s2, l::add);",
+            "        }",
+            "}")
+        .doTest();
+  }
+
   private CompilationTestHelper makeHelper() {
     return makeTestHelperWithArgs(
         JSpecifyJavacConfig.withJSpecifyModeArgs(

test-java-lib/src/main/java/com/uber/lib/unannotated/ProviderNullMarkedViaModel.java

@@ -2,4 +2,8 @@
 
 public interface ProviderNullMarkedViaModel<T> {
   T get();
+
+  static <U> ProviderNullMarkedViaModel<U> of(U value) {
+    return () -> value;
+  }
 }

test-library-models/src/main/java/com/uber/nullaway/testlibrarymodels/TestLibraryModels.java

@@ -143,4 +143,10 @@ public ImmutableSetMultimap<String, Integer> typeVariablesWithNullableUpperBound
   public ImmutableSet<String> nullMarkedClasses() {
     return ImmutableSet.of("com.uber.lib.unannotated.ProviderNullMarkedViaModel");
   }
+
+  @Override
+  public ImmutableSetMultimap<MethodRef, Integer> methodTypeVariablesWithNullableUpperBounds() {
+    return ImmutableSetMultimap.of(
+        methodRef("com.uber.lib.unannotated.ProviderNullMarkedViaModel", "<U>of(U)"), 0);
+  }
 }