JSpecify: handle Nullability for lambda expression parameters for Gen…

…eric Types (#852)

We were not getting the desired error for the positive test in the
following test case:

```java
class Test {
  interface A<T1 extends @nullable Object> {
    String function(T1 o);
  }
  static void testPositive() {
        // We were not getting the desired error here and these changes address that
	// BUG: Diagnostic contains: dereferenced expression o is @nullable
	A<@nullable Object> p = o -> o.toString();
  }
  static void testNegative() {
    A<Object> p = o -> o.toString();
  }
}
```

Upon changing the way we compute the Nullability of the functional
interface parameters, we are able to get the desired error.

Note: These changes only affect lambda parameters and do not deal with
return types. That will be addressed in a future PR.

---------

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

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

@@ -3,10 +3,12 @@
 import static com.uber.nullaway.Nullness.NONNULL;
 import static com.uber.nullaway.Nullness.NULLABLE;
 
+import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.ClassTree;
 import com.sun.source.tree.LambdaExpressionTree;
 import com.sun.source.tree.VariableTree;
 import com.sun.tools.javac.code.Symbol;
+import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.Types;
 import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.CodeAnnotationInfo;
@@ -92,13 +94,27 @@ private static NullnessStore lambdaInitialStore(
     Symbol.MethodSymbol fiMethodSymbol = NullabilityUtil.getFunctionalInterfaceMethod(code, types);
     com.sun.tools.javac.util.List<Symbol.VarSymbol> fiMethodParameters =
         fiMethodSymbol.getParameters();
+    // This obtains the types of the functional interface method parameters with preserved
+    // annotations in case of generic type arguments.  Only used in JSpecify mode.
+    List<Type> overridenMethodParamTypeList =
+        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);
     if (isFIAnnotated) {
       for (int i = 0; i < fiMethodParameters.size(); i++) {
-        fiArgumentPositionNullness[i] =
-            Nullness.hasNullableAnnotation(fiMethodParameters.get(i), config) ? NULLABLE : NONNULL;
+        if (Nullness.hasNullableAnnotation(fiMethodParameters.get(i), config)) {
+          // Get the Nullness if the Annotation is directly written with the parameter
+          fiArgumentPositionNullness[i] = NULLABLE;
+        } else if (config.isJSpecifyMode()
+            && Nullness.hasNullableAnnotation(
+                overridenMethodParamTypeList.get(i).getAnnotationMirrors().stream(), config)) {
+          // Get the Nullness if the Annotation is indirectly applied through a generic type if we
+          // are in JSpecify mode
+          fiArgumentPositionNullness[i] = NULLABLE;
+        } else {
+          fiArgumentPositionNullness[i] = NONNULL;
+        }
       }
     }
     fiArgumentPositionNullness =

nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java

@@ -585,7 +585,7 @@ public void testForMethodReferenceAsMethodParameter() {
   }
 
   @Test
-  public void testForLambdasInAnAssignment() {
+  public void testForLambdasInAnAssignmentWithSingleParam() {
     makeHelper()
         .addSourceLines(
             "Test.java",
@@ -596,8 +596,51 @@ public void testForLambdasInAnAssignment() {
             "    String function(T1 o);",
             "  }",
             "  static void testPositive() {",
-            "    // TODO: we should report an error here, since the lambda cannot take",
-            "    // a @Nullable parameter.  we don't catch this yet",
+            "    // BUG: Diagnostic contains: dereferenced expression o is @Nullable",
+            "    A<@Nullable Object> p = o -> o.toString();",
+            "  }",
+            "  static void testNegative() {",
+            "    A<Object> p = o -> o.toString();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void testForLambdasInAnAssignmentWithMultipleParams() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "class Test {",
+            "  interface A<T1 extends @Nullable Object,T2 extends @Nullable Object> {",
+            "    String function(T1 o1,T2 o2);",
+            "  }",
+            "  static void testPositive() {",
+            "    // BUG: Diagnostic contains: dereferenced expression o1 is @Nullable",
+            "    A<@Nullable Object,Object> p = (o1,o2) -> o1.toString();",
+            "  }",
+            "  static void testNegative() {",
+            "    A<@Nullable Object,Object> p = (o1,o2) -> o2.toString();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void testForLambdasInAnAssignmentWithoutJSpecifyMode() {
+    makeHelperWithoutJSpecifyMode()
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "class Test {",
+            "  interface A<T1 extends @Nullable Object> {",
+            "    String function(T1 o);",
+            "  }",
+            "  static void testPositive() {",
+            "    // Using outside JSpecify Mode So we don't get a bug here",
             "    A<@Nullable Object> p = o -> o.toString();",
             "  }",
             "  static void testNegative() {",
@@ -1427,4 +1470,8 @@ private CompilationTestHelper makeHelper() {
         Arrays.asList(
             "-XepOpt:NullAway:AnnotatedPackages=com.uber", "-XepOpt:NullAway:JSpecifyMode=true"));
   }
+
+  private CompilationTestHelper makeHelperWithoutJSpecifyMode() {
+    return makeTestHelperWithArgs(Arrays.asList("-XepOpt:NullAway:AnnotatedPackages=com.uber"));
+  }
 }