Adding support for nullable type upper bounds considering Library mod…

…els (#903)

The below line previously reported an error considering the upperbound
of the generic type parameter couldn't be Nullable, adding it into
library models to allow it to be Nullable passes the test case.

```java
Function<String,@nullable String> removeA = a -> a.replace("a","");
```

We now have library model methods for treating a type as `@NullMarked`
and also for making the upper bound of type variables `@Nullable`. We
use these new methods to model `java.util.function.Function` as
`@NullMarked`, in JSpecify mode only.

We fix a related issue with unwanted errors when passing a lambda /
method reference to `@NullUnmarked` code; this was necessary to get
NullAway itself to still build with NullAway checking enabled.

(Due to issues with the branch this PR was recreated from the original
#893)

---------

Co-authored-by: Manu Sridharan <msridhar@gmail.com>

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

@@ -32,8 +32,10 @@
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.util.Context;
+import com.uber.nullaway.handlers.Handler;
 import java.util.HashMap;
 import java.util.Map;
+import javax.annotation.Nullable;
 import javax.lang.model.element.ElementKind;
 
 /**
@@ -120,7 +122,7 @@ public boolean isGenerated(Symbol symbol, Config config) {
               symbol));
       return false;
     }
-    Symbol.ClassSymbol outermostClassSymbol = get(classSymbol, config).outermostClassSymbol;
+    Symbol.ClassSymbol outermostClassSymbol = get(classSymbol, config, null).outermostClassSymbol;
     return hasDirectAnnotationWithSimpleName(outermostClassSymbol, "Generated");
   }
 
@@ -145,10 +147,11 @@ private static boolean isClassFieldOfPrimitiveType(Symbol symbol) {
    *
    * @param symbol symbol for entity
    * @param config NullAway config
+   * @param handler optional NullAway Handler
    * @return true if symbol represents an entity contained in a class that is unannotated; false
    *     otherwise
    */
-  public boolean isSymbolUnannotated(Symbol symbol, Config config) {
+  public boolean isSymbolUnannotated(Symbol symbol, Config config, @Nullable Handler handler) {
     Symbol.ClassSymbol classSymbol;
     if (symbol instanceof Symbol.ClassSymbol) {
       classSymbol = (Symbol.ClassSymbol) symbol;
@@ -162,7 +165,7 @@ public boolean isSymbolUnannotated(Symbol symbol, Config config) {
     } else {
       classSymbol = castToNonNull(ASTHelpers.enclosingClass(symbol));
     }
-    final ClassCacheRecord classCacheRecord = get(classSymbol, config);
+    final ClassCacheRecord classCacheRecord = get(classSymbol, config, handler);
     boolean inAnnotatedClass = classCacheRecord.isNullnessAnnotated;
     if (symbol.getKind().equals(ElementKind.METHOD)
         || symbol.getKind().equals(ElementKind.CONSTRUCTOR)) {
@@ -176,12 +179,15 @@ public boolean isSymbolUnannotated(Symbol symbol, Config config) {
    * Check whether a class should be treated as nullness-annotated.
    *
    * @param classSymbol The symbol for the class to be checked
+   * @param config NullAway config
+   * @param handler optional NullAway handler
    * @return Whether this class should be treated as null-annotated, taking into account annotations
    *     on enclosing classes, the containing package, and other NullAway configuration like
    *     annotated packages
    */
-  public boolean isClassNullAnnotated(Symbol.ClassSymbol classSymbol, Config config) {
-    return get(classSymbol, config).isNullnessAnnotated;
+  public boolean isClassNullAnnotated(
+      Symbol.ClassSymbol classSymbol, Config config, Handler handler) {
+    return get(classSymbol, config, handler).isNullnessAnnotated;
   }
 
   /**
@@ -190,12 +196,15 @@ public boolean isClassNullAnnotated(Symbol.ClassSymbol classSymbol, Config confi
    * <p>This method is recursive, using the cache on the way up and populating it on the way down.
    *
    * @param classSymbol The class to query, possibly an inner class
+   * @param config NullAway config
+   * @param handler optional NullAway handler
    * @return A record including the outermost class in which the given class is nested, as well as
    *     boolean flag noting whether it should be treated as nullness-annotated, taking into account
    *     annotations on enclosing classes, the containing package, and other NullAway configuration
    *     like annotated packages
    */
-  private ClassCacheRecord get(Symbol.ClassSymbol classSymbol, Config config) {
+  private ClassCacheRecord get(
+      Symbol.ClassSymbol classSymbol, Config config, @Nullable Handler handler) {
     ClassCacheRecord record = classCache.getIfPresent(classSymbol);
     if (record != null) {
       return record;
@@ -211,7 +220,7 @@ private ClassCacheRecord get(Symbol.ClassSymbol classSymbol, Config config) {
       Symbol.ClassSymbol enclosingClass = ASTHelpers.enclosingClass(classSymbol);
       // enclosingClass can be null in weird cases like for array methods
       if (enclosingClass != null) {
-        ClassCacheRecord recordForEnclosing = get(enclosingClass, config);
+        ClassCacheRecord recordForEnclosing = get(enclosingClass, config, handler);
         // Check if this class is annotated, recall that enclosed scopes override enclosing scopes
         boolean isAnnotated = recordForEnclosing.isNullnessAnnotated;
         if (enclosingMethod != null) {
@@ -232,7 +241,8 @@ record = new ClassCacheRecord(recordForEnclosing.outermostClassSymbol, isAnnotat
     }
     if (record == null) {
       // We are already at the outermost class (we can find), so let's create a record for it
-      record = new ClassCacheRecord(classSymbol, isAnnotatedTopLevelClass(classSymbol, config));
+      record =
+          new ClassCacheRecord(classSymbol, isAnnotatedTopLevelClass(classSymbol, config, handler));
     }
     classCache.put(classSymbol, record);
     return record;
@@ -258,7 +268,17 @@ private boolean shouldTreatAsUnannotated(Symbol.ClassSymbol classSymbol, Config
     return false;
   }
 
-  private boolean isAnnotatedTopLevelClass(Symbol.ClassSymbol classSymbol, Config config) {
+  /**
+   * Check if a non-nested class should be treated as nullness-annotated.
+   *
+   * @param classSymbol symbol for the class
+   * @param config NullAway config
+   * @param handler optional NullAway handler. If present, used to check if a class is NullMarked
+   *     based on {@link Handler#onOverrideNullMarkedClasses(String)}
+   * @return {@code true} iff the class should be treated as nullness-annotated
+   */
+  private boolean isAnnotatedTopLevelClass(
+      Symbol.ClassSymbol classSymbol, Config config, @Nullable Handler handler) {
     // First, check for an explicitly @NullUnmarked top level class
     if (hasDirectAnnotationWithSimpleName(classSymbol, NullabilityUtil.NULLUNMARKED_SIMPLE_NAME)) {
       return false;
@@ -269,7 +289,12 @@ private boolean isAnnotatedTopLevelClass(Symbol.ClassSymbol classSymbol, Config
       // make sure it's not explicitly configured as unannotated
       return !shouldTreatAsUnannotated(classSymbol, config);
     }
-    return false;
+    // Check if it is NullMarked inside a Library Model when in JSpecify Mode
+    if (config.isJSpecifyMode() && handler != null) {
+      return handler.onOverrideNullMarkedClasses(classSymbol.toString());
+    } else {
+      return false;
+    }
   }
 
   /**

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

@@ -112,6 +112,26 @@ public interface LibraryModels {
    */
   ImmutableSet<MethodRef> nonNullReturns();
 
+  /**
+   * 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.
+   *
+   * @return map from the className to the positions of the generic type arguments that have a
+   *     {@code Nullable} upper bound.
+   */
+  default ImmutableSetMultimap<String, Integer> typeVariablesWithNullableUpperBounds() {
+    return ImmutableSetMultimap.of();
+  }
+
+  /**
+   * Get the set of library classes that are NullMarked. Only used in JSpecify mode.
+   *
+   * @return set of library classes that are NullMarked.
+   */
+  default ImmutableSet<String> nullMarkedClasses() {
+    return ImmutableSet.of();
+  }
+
   /**
    * Get (method, parameter) pairs that act as castToNonNull(...) methods.
    *
@@ -137,7 +157,9 @@ public interface LibraryModels {
    *
    * @return set of library fields that may be {@code null}.
    */
-  ImmutableSet<FieldRef> nullableFields();
+  default ImmutableSet<FieldRef> nullableFields() {
+    return ImmutableSet.of();
+  }
 
   /**
    * Get a list of custom stream library specifications.

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

@@ -307,7 +307,7 @@ public NullAway(ErrorProneFlags flags) {
   private boolean isMethodUnannotated(MethodInvocationNode invocationNode) {
     return invocationNode == null
         || codeAnnotationInfo.isSymbolUnannotated(
-            ASTHelpers.getSymbol(invocationNode.getTree()), config);
+            ASTHelpers.getSymbol(invocationNode.getTree()), config, handler);
   }
 
   private boolean withinAnnotatedCode(VisitorState state) {
@@ -346,7 +346,7 @@ private boolean checkMarkingForPath(VisitorState state) {
     if (enclosingMarkableSymbol == null) {
       return false;
     }
-    return !codeAnnotationInfo.isSymbolUnannotated(enclosingMarkableSymbol, config);
+    return !codeAnnotationInfo.isSymbolUnannotated(enclosingMarkableSymbol, config, handler);
   }
 
   @Override
@@ -612,7 +612,7 @@ private void checkForMethodNullMarkedness(MethodTree tree, VisitorState state) {
             methodSymbol, NullabilityUtil.NULLMARKED_SIMPLE_NAME)) {
           // We still care here if this is a transition between @NullUnmarked and @NullMarked code,
           // within partially marked code, see checks below for markedMethodInUnmarkedContext.
-          if (!codeAnnotationInfo.isClassNullAnnotated(methodSymbol.enclClass(), config)) {
+          if (!codeAnnotationInfo.isClassNullAnnotated(methodSymbol.enclClass(), config, handler)) {
             markedMethodInUnmarkedContext = true;
           }
         }
@@ -713,7 +713,8 @@ public Description matchParameterizedType(ParameterizedTypeTree tree, VisitorSta
       return Description.NO_MATCH;
     }
     if (config.isJSpecifyMode()) {
-      GenericsChecks.checkInstantiationForParameterizedTypedTree(tree, state, this, config);
+      GenericsChecks.checkInstantiationForParameterizedTypedTree(
+          tree, state, this, config, handler);
     }
     return Description.NO_MATCH;
   }
@@ -741,7 +742,7 @@ private Description checkParamOverriding(
         (memberReferenceTree != null)
             && ((JCTree.JCMemberReference) memberReferenceTree).kind.isUnbound();
     final boolean isOverriddenMethodAnnotated =
-        !codeAnnotationInfo.isSymbolUnannotated(overriddenMethod, config);
+        !codeAnnotationInfo.isSymbolUnannotated(overriddenMethod, config, handler);
 
     // Get argument nullability for the overridden method.  If overriddenMethodArgNullnessMap[i] is
     // null, parameter i is treated as unannotated.
@@ -861,7 +862,8 @@ static Trees getTreesInstance(VisitorState state) {
 
   private Nullness getMethodReturnNullness(
       Symbol.MethodSymbol methodSymbol, VisitorState state, Nullness defaultForUnannotated) {
-    final boolean isMethodAnnotated = !codeAnnotationInfo.isSymbolUnannotated(methodSymbol, config);
+    final boolean isMethodAnnotated =
+        !codeAnnotationInfo.isSymbolUnannotated(methodSymbol, config, handler);
     Nullness methodReturnNullness =
         defaultForUnannotated; // Permissive default for unannotated code.
     if (isMethodAnnotated) {
@@ -916,13 +918,15 @@ private Description checkReturnExpression(
     // type is @Nullable, and if so, bail out.
     if (getMethodReturnNullness(methodSymbol, state, Nullness.NULLABLE).equals(Nullness.NULLABLE)) {
       return Description.NO_MATCH;
-    } else if (config.isJSpecifyMode()
-        && lambdaTree != null
-        && GenericsChecks.getGenericMethodReturnTypeNullness(
-                methodSymbol, ASTHelpers.getType(lambdaTree), state, config)
-            .equals(Nullness.NULLABLE)) {
-      // In JSpecify mode, the return type of a lambda may be @Nullable via a type argument
-      return Description.NO_MATCH;
+    } else if (config.isJSpecifyMode() && lambdaTree != null) {
+      if (GenericsChecks.getGenericMethodReturnTypeNullness(
+                  methodSymbol, ASTHelpers.getType(lambdaTree), state, config)
+              .equals(Nullness.NULLABLE)
+          || GenericsChecks.passingLambdaOrMethodRefWithGenericReturnToUnmarkedCode(
+              methodSymbol, lambdaTree, state, config, codeAnnotationInfo, handler)) {
+        // In JSpecify mode, the return type of a lambda may be @Nullable via a type argument
+        return Description.NO_MATCH;
+      }
     }
 
     // Return type is @NonNull.  Check if the expression is @Nullable
@@ -950,7 +954,7 @@ public Description matchLambdaExpression(LambdaExpressionTree tree, VisitorState
     // (like Rx nullability) run dataflow analysis
     updateEnvironmentMapping(state.getPath(), state);
     handler.onMatchLambdaExpression(this, tree, state, funcInterfaceMethod);
-    if (codeAnnotationInfo.isSymbolUnannotated(funcInterfaceMethod, config)) {
+    if (codeAnnotationInfo.isSymbolUnannotated(funcInterfaceMethod, config, handler)) {
       return Description.NO_MATCH;
     }
     Description description =
@@ -1065,8 +1069,10 @@ private boolean overriddenMethodReturnsNonNull(
         // For a method reference, we get generic type arguments from javac's inferred type for the
         // tree, which properly preserves type-use annotations
         return GenericsChecks.getGenericMethodReturnTypeNullness(
-                overriddenMethod, ASTHelpers.getType(memberReferenceTree), state, config)
-            .equals(Nullness.NONNULL);
+                    overriddenMethod, ASTHelpers.getType(memberReferenceTree), state, config)
+                .equals(Nullness.NONNULL)
+            && !GenericsChecks.passingLambdaOrMethodRefWithGenericReturnToUnmarkedCode(
+                overriddenMethod, memberReferenceTree, state, config, codeAnnotationInfo, handler);
       } else {
         // Use the enclosing class of the overriding method to find generic type arguments
         return GenericsChecks.getGenericMethodReturnTypeNullness(
@@ -1176,7 +1182,7 @@ private boolean relevantInitializerMethodOrBlock(
       // in those cases, we won't even have a populated class2Entities map). We skip this check if
       // we are not inside a @NullMarked/annotated *class*:
       if (nullMarkingForTopLevelClass == NullMarking.PARTIALLY_MARKED
-          && !codeAnnotationInfo.isClassNullAnnotated(enclClassSymbol, config)) {
+          && !codeAnnotationInfo.isClassNullAnnotated(enclClassSymbol, config, handler)) {
         return false;
       }
 
@@ -1697,7 +1703,8 @@ private Description handleInvocation(
       return Description.NO_MATCH;
     }
 
-    final boolean isMethodAnnotated = !codeAnnotationInfo.isSymbolUnannotated(methodSymbol, config);
+    final boolean isMethodAnnotated =
+        !codeAnnotationInfo.isSymbolUnannotated(methodSymbol, config, handler);
     // If argumentPositionNullness[i] == null, parameter i is unannotated
     Nullness[] argumentPositionNullness = new Nullness[formalParams.size()];
 
@@ -2287,7 +2294,7 @@ private boolean isExcludedClass(Symbol.ClassSymbol classSymbol) {
     if (config.isExcludedClass(className)) {
       return true;
     }
-    if (!codeAnnotationInfo.isClassNullAnnotated(classSymbol, config)) {
+    if (!codeAnnotationInfo.isClassNullAnnotated(classSymbol, config, handler)) {
       return true;
     }
     // check annotations
@@ -2431,7 +2438,7 @@ private boolean mayBeNullExpr(VisitorState state, ExpressionTree expr) {
 
   private boolean mayBeNullMethodCall(
       Symbol.MethodSymbol exprSymbol, MethodInvocationTree invocationTree, VisitorState state) {
-    if (codeAnnotationInfo.isSymbolUnannotated(exprSymbol, config)) {
+    if (codeAnnotationInfo.isSymbolUnannotated(exprSymbol, config, handler)) {
       return false;
     }
     if (Nullness.hasNullableAnnotation(exprSymbol, config)) {

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

@@ -362,7 +362,7 @@ public static boolean mayBeNullFieldFromType(
       Symbol symbol, Config config, CodeAnnotationInfo codeAnnotationInfo) {
     return !(symbol.getSimpleName().toString().equals("class")
             || symbol.isEnum()
-            || codeAnnotationInfo.isSymbolUnannotated(symbol, config))
+            || codeAnnotationInfo.isSymbolUnannotated(symbol, config, null))
         && Nullness.hasNullableAnnotation(symbol, config);
   }
 

nullaway/src/main/java/com/uber/nullaway/dataflow/CoreNullnessStoreInitializer.java

@@ -100,7 +100,8 @@ private static NullnessStore lambdaInitialStore(
         types.memberType(ASTHelpers.getType(code), fiMethodSymbol).getParameterTypes();
     // If fiArgumentPositionNullness[i] == null, parameter position i is unannotated
     Nullness[] fiArgumentPositionNullness = new Nullness[fiMethodParameters.size()];
-    final boolean isFIAnnotated = !codeAnnotationInfo.isSymbolUnannotated(fiMethodSymbol, config);
+    final boolean isFIAnnotated =
+        !codeAnnotationInfo.isSymbolUnannotated(fiMethodSymbol, config, handler);
     if (isFIAnnotated) {
       for (int i = 0; i < fiMethodParameters.size(); i++) {
         if (Nullness.hasNullableAnnotation(fiMethodParameters.get(i), config)) {