Use ground target types when handling lambdas and method refs passed to generic methods

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

@@ -829,7 +829,7 @@ public void checkTypeParameterNullnessForAssignability(Tree tree, VisitorState s
     if (!assignedToLocal
         && (rhsTree instanceof LambdaExpressionTree || rhsTree instanceof MemberReferenceTree)
         && isAssignmentToField(tree)) {
-      maybeStorePolyExpressionTypeFromTarget(rhsTree, lhsType);
+      maybeStorePolyExpressionTypeFromTarget(rhsTree, lhsType, state);
     }
     boolean varLocalDeclaration =
         tree instanceof VariableTree varTree && isVarLocalVariableDeclaration(varTree);
@@ -1081,9 +1081,15 @@ private MethodInferenceResult runInferenceForCall(
                       || argument instanceof MemberReferenceTree) {
                     Type polyExprTreeType = ASTHelpers.getType(argument);
                     if (polyExprTreeType != null) {
+                      Type formalParamGroundTargetType =
+                          GenericsUtils.groundTargetType(formalParamType, state, config);
                       Type typeWithInferredNullability =
                           TypeSubstitutionUtils.updateTypeWithInferredNullability(
-                              polyExprTreeType, formalParamType, typeVarNullability, state, config);
+                              polyExprTreeType,
+                              formalParamGroundTargetType,
+                              typeVarNullability,
+                              state,
+                              config);
                       inferredPolyExpressionTypes.put(argument, typeWithInferredNullability);
                     }
                   }
@@ -1255,10 +1261,11 @@ private void handleLambdaInGenericMethodInference(
     Symbol.MethodSymbol fiMethod =
         NullabilityUtil.getFunctionalInterfaceMethod(lambda, state.getTypes());
 
+    Type groundTargetType = GenericsUtils.groundTargetType(lhsType, state, config);
     // get the return type of the functional interface method, viewed as a member of the lhs
     // type, so the generic method's type variables are substituted in
     Type.MethodType fiMethodTypeAsMember =
-        TypeSubstitutionUtils.memberType(state.getTypes(), lhsType, fiMethod, config)
+        TypeSubstitutionUtils.memberType(state.getTypes(), groundTargetType, fiMethod, config)
             .asMethodType();
     Type fiReturnType = fiMethodTypeAsMember.getReturnType();
     Tree body = lambda.getBody();
@@ -1308,9 +1315,10 @@ private void handleMethodRefInGenericMethodInference(
       ConstraintSolver solver,
       Type lhsType,
       MemberReferenceTree memberReferenceTree) {
+    Type groundTargetType = GenericsUtils.groundTargetType(lhsType, state, config);
     GenericsUtils.processMethodRefTypeRelations(
         this,
-        lhsType,
+        groundTargetType,
         memberReferenceTree,
         state,
         (subtype, supertype, unused) -> {
@@ -1799,12 +1807,14 @@ public void compareGenericTypeParameterNullabilityForCall(
               }
 
               if (currentActualParam instanceof MemberReferenceTree memberReferenceTree) {
+                Type groundFormalParameter =
+                    GenericsUtils.groundTargetType(formalParameter, state, config);
                 // the type of the method reference tree provided by javac may not capture
                 // nullability of nested types. So, do explicit type checks based on the return and
                 // parameter types of the referenced method
                 GenericsUtils.processMethodRefTypeRelations(
                     this,
-                    formalParameter,
+                    groundFormalParameter,
                     memberReferenceTree,
                     state,
                     (subtype, supertype, relationKind) -> {
@@ -1818,14 +1828,14 @@ public void compareGenericTypeParameterNullabilityForCall(
                         }
                       }
                     });
-                maybeStorePolyExpressionTypeFromTarget(currentActualParam, formalParameter);
+                maybeStorePolyExpressionTypeFromTarget(currentActualParam, formalParameter, state);
                 return;
               }
 
               TreePath pathToParam = pathWithLeaf(state.getPath(), currentActualParam);
               Type actualParameterType;
               if (currentActualParam instanceof LambdaExpressionTree) {
-                maybeStorePolyExpressionTypeFromTarget(currentActualParam, formalParameter);
+                maybeStorePolyExpressionTypeFromTarget(currentActualParam, formalParameter, state);
               }
               Type inferredPolyType = inferredPolyExpressionTypes.get(currentActualParam);
               if (inferredPolyType != null) {
@@ -2683,17 +2693,19 @@ public boolean isNullableAnnotated(Type type) {
    * <p>This is used to compensate for javac dropping annotations on type variables in poly
    * expression target types, so later checks use the correctly annotated functional interface type.
    */
-  private void maybeStorePolyExpressionTypeFromTarget(Tree polyExpressionTree, Type targetType) {
+  private void maybeStorePolyExpressionTypeFromTarget(
+      Tree polyExpressionTree, Type targetType, VisitorState state) {
     if (targetType.isRaw() || inferredPolyExpressionTypes.containsKey(polyExpressionTree)) {
       return;
     }
     Type polyExpressionType = ASTHelpers.getType(polyExpressionTree);
     if (polyExpressionType == null) {
       return;
     }
+    Type groundTargetType = GenericsUtils.groundTargetType(targetType, state, config);
     Type polyExpressionTypeWithTargetAnnotations =
         TypeSubstitutionUtils.restoreExplicitNullabilityAnnotations(
-            targetType, polyExpressionType, config, Collections.emptyMap());
+            groundTargetType, polyExpressionType, config, Collections.emptyMap());
     inferredPolyExpressionTypes.put(polyExpressionTree, polyExpressionTypeWithTargetAnnotations);
   }
 

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

@@ -1,5 +1,7 @@
 package com.uber.nullaway.generics;
 
+import static com.uber.nullaway.NullabilityUtil.castToNonNull;
+
 import com.google.common.base.Verify;
 import com.google.errorprone.VisitorState;
 import com.google.errorprone.util.ASTHelpers;
@@ -11,9 +13,13 @@
 import com.sun.tools.javac.code.Symtab;
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.Type.CapturedType;
+import com.sun.tools.javac.code.Type.ClassType;
 import com.sun.tools.javac.code.Type.WildcardType;
 import com.sun.tools.javac.code.Types;
 import com.sun.tools.javac.tree.JCTree;
+import com.sun.tools.javac.util.List;
+import com.sun.tools.javac.util.ListBuffer;
+import com.uber.nullaway.Config;
 import com.uber.nullaway.NullabilityUtil;
 import javax.lang.model.type.TypeKind;
 import org.jspecify.annotations.Nullable;
@@ -78,6 +84,60 @@ static Type wildcardUpperBound(WildcardType wildcardType, VisitorState state) {
     return null;
   }
 
+  /**
+   * Returns a non-wildcard functional interface parameterization for lambda and method-reference
+   * checking. For immediate wildcard type arguments, use the bound that determines the functional
+   * interface descriptor, preserving wildcards in nested type positions.
+   *
+   * <p>This implements the ground target type behavior used for lambda and method-reference target
+   * typing; see JLS <a
+   * href="https://docs.oracle.com/javase/specs/jls/se21/html/jls-15.html#jls-15.27.3">15.27.3</a>,
+   * JLS <a
+   * href="https://docs.oracle.com/javase/specs/jls/se21/html/jls-15.html#jls-15.13.2">15.13.2</a>,
+   * and the non-wildcard parameterization rules in JLS <a
+   * href="https://docs.oracle.com/javase/specs/jls/se21/html/jls-9.html#jls-9.9">9.9</a>.
+   */
+  @SuppressWarnings("ReferenceEquality")
+  static Type groundTargetType(Type targetType, VisitorState state, Config config) {
+    if (!config.handleWildcardGenerics()) {
+      return targetType;
+    }
+    if (!(targetType instanceof ClassType classType) || targetType.isRaw()) {
+      return targetType;
+    }
+    List<Type> typeArguments = classType.getTypeArguments();
+    if (typeArguments.isEmpty()) {
+      return targetType;
+    }
+    ListBuffer<Type> groundedTypeArguments = new ListBuffer<>();
+    boolean changed = false;
+    for (Type typeArgument : typeArguments) {
+      Type groundedTypeArgument = groundTypeArgument(typeArgument, state);
+      groundedTypeArguments.append(groundedTypeArgument);
+      changed |= groundedTypeArgument != typeArgument;
+    }
+    return changed
+        ? TypeMetadataBuilder.TYPE_METADATA_BUILDER.createClassType(
+            targetType, classType.getEnclosingType(), groundedTypeArguments.toList())
+        : targetType;
+  }
+
+  /**
+   * Grounds one immediate wildcard type argument according to the non-wildcard parameterization
+   * rules for functional interface target types in JLS <a
+   * href="https://docs.oracle.com/javase/specs/jls/se21/html/jls-9.html#jls-9.9">9.9</a>.
+   */
+  private static Type groundTypeArgument(Type typeArgument, VisitorState state) {
+    WildcardType wildcardType = asWildcard(typeArgument);
+    if (wildcardType == null) {
+      return typeArgument;
+    }
+    if (wildcardType.kind == BoundKind.SUPER) {
+      return castToNonNull(wildcardType.getSuperBound());
+    }
+    return wildcardUpperBound(wildcardType, state);
+  }
+
   /**
    * Handler for method reference type relations, used by {{@link
    * #processMethodRefTypeRelations(GenericsChecks, Type, MemberReferenceTree, VisitorState,

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

@@ -4,7 +4,6 @@
 import com.uber.nullaway.NullAwayTestsBase;
 import com.uber.nullaway.generics.JSpecifyJavacConfig;
 import java.util.Arrays;
-import org.junit.Ignore;
 import org.junit.Test;
 
 public class WildcardTests extends NullAwayTestsBase {
@@ -580,7 +579,6 @@ static void test() {
         .doTest();
   }
 
-  @Ignore("https://github.com/uber/NullAway/issues/1522")
   @Test
   public void issue1522() {
     makeHelperWithInferenceFailureWarning()
@@ -605,6 +603,163 @@ static <T> Foo<T> after(Foo<Optional<T>> foo) {
         .doTest();
   }
 
+  @Test
+  public void issue1522SelfContained() {
+    makeHelperWithInferenceFailureWarning()
+        .addSourceLines(
+            "Test.java",
+            """
+            import org.jspecify.annotations.*;
+            @NullMarked
+            class Test {
+              interface Function<T extends @Nullable Object, U extends @Nullable Object> {
+                U apply(T t);
+              }
+              static class Optional<T> {
+                public @Nullable T orElse(@Nullable T other) {
+                    throw new RuntimeException();
+                }
+              }
+              static class Foo<T> {
+                public final <V> Foo<V> mapNotNull(Function<? super T, ? extends @Nullable V> mapper) {
+                  throw new RuntimeException();
+                }
+              }
+              static <T> Foo<T> after(Foo<Optional<T>> foo) {
+                return foo.mapNotNull(x -> x.orElse(null));
+              }
+            }
+            """)
+        .doTest();
+  }
+
+  @Test
+  public void issue1522SelfContainedWithMethodReference() {
+    makeHelperWithInferenceFailureWarning()
+        .addSourceLines(
+            "Test.java",
+            """
+            import org.jspecify.annotations.*;
+            @NullMarked
+            class Test {
+              interface Function<T extends @Nullable Object, U extends @Nullable Object> {
+                U apply(T t);
+              }
+              static class Optional<T> {
+                public @Nullable T orElse(@Nullable T other) {
+                    throw new RuntimeException();
+                }
+              }
+              static class Foo<T> {
+                public final <V> Foo<V> mapNotNull(Function<? super T, ? extends @Nullable V> mapper) {
+                  throw new RuntimeException();
+                }
+              }
+              static <T> @Nullable T orElseNull(Optional<T> optional) {
+                return optional.orElse(null);
+              }
+              static <T> Foo<T> after(Foo<Optional<T>> foo) {
+                return foo.mapNotNull(Test::orElseNull);
+              }
+            }
+            """)
+        .doTest();
+  }
+
+  @Test
+  public void groundTargetTypePreservesNestedWildcards() {
+    makeHelperWithInferenceFailureWarning()
+        .addSourceLines(
+            "Test.java",
+            """
+            import org.jspecify.annotations.*;
+            @NullMarked
+            class Test {
+              interface Function<T extends @Nullable Object, R extends @Nullable Object> {
+                R apply(T t);
+              }
+              static class Box<T extends @Nullable Object> {
+                T get() {
+                  throw new RuntimeException();
+                }
+              }
+              static <R extends @Nullable Object> R invokeNested(
+                  Function<Box<? super String>, R> mapper) {
+                throw new RuntimeException();
+              }
+              static <R extends @Nullable Object> R invokeNestedWithUpperBound(
+                  Function<Box<? extends String>, R> mapper) {
+                throw new RuntimeException();
+              }
+              static <R extends @Nullable Object> R invokeTopLevelWildcard(
+                  Function<? super Box<? super String>, R> mapper) {
+                throw new RuntimeException();
+              }
+              static <R extends @Nullable Object> R invokeArray(
+                  Function<Box<? super String>[], R> mapper) {
+                throw new RuntimeException();
+              }
+              static void testNestedWildcard() {
+                invokeNested(box -> {
+                  // BUG: Diagnostic contains: dereferenced expression box.get() is @Nullable
+                  box.get().hashCode();
+                  return null;
+                });
+                invokeNestedWithUpperBound(box -> {
+                  // safe since the upper bound of the Box type variable is @NonNull String,
+                  // so box.get() cannot be null
+                  box.get().hashCode();
+                  return null;
+                });
+              }
+              static void testTopLevelWildcardBound() {
+                invokeTopLevelWildcard(box -> {
+                  // BUG: Diagnostic contains: dereferenced expression box.get() is @Nullable
+                  box.get().hashCode();
+                  return null;
+                });
+              }
+              static void testArrayWithNestedWildcard() {
+                invokeArray(boxes -> {
+                  // BUG: Diagnostic contains: dereferenced expression boxes[0].get() is @Nullable
+                  boxes[0].get().hashCode();
+                  return null;
+                });
+              }
+            }
+            """)
+        .doTest();
+  }
+
+  @Test
+  public void groundTargetTypePreservesNestedWildcardsForMethodReferences() {
+    makeHelperWithInferenceFailureWarning()
+        .addSourceLines(
+            "Test.java",
+            """
+            import org.jspecify.annotations.*;
+            @NullMarked
+            class Test {
+              interface Function<T extends @Nullable Object, R extends @Nullable Object> {
+                R apply(T t);
+              }
+              static class Box<T extends @Nullable Object> {}
+              static <R extends @Nullable Object> R invokeExtendsNullable(
+                  Function<Box<? extends @Nullable String>, R> mapper) {
+                throw new RuntimeException();
+              }
+              static @Nullable Object needsBoxExtendsString(Box<? extends String> box) {
+                return null;
+              }
+              static void test() {
+                // BUG: Diagnostic contains: parameter type of referenced method is Box<? extends String>
+                invokeExtendsNullable(Test::needsBoxExtendsString);
+              }
+            }
+            """)
+        .doTest();
+  }
+
   /**
    * Extracted from Caffeine; exposed some subtle bugs in substitutions involving identity of {@code
    * Type} objects