JSpecify: properly handle lambdas in anonymous inner classes (#1165)

Fixes #1156

When checking correctness of lambda overrides we get the type from
`javac`'s inferred type rather than using the enclosing class. Also,
when getting the type of an anonymous class we add some extra assertions
to ensure we find the right `NewClassTree` representing the anonymous
class. I _think_ we'll always get the right tree now, but the asserts
give peace of mind.

Author: msridhar

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

@@ -807,12 +807,17 @@ private Description checkParamOverriding(
         } else if (config.isJSpecifyMode()) {
           // Check if the parameter type is a type variable and the corresponding generic type
           // argument is @Nullable
-          if (memberReferenceTree != null) {
-            // For a method reference, we get generic type arguments from the javac's inferred type
-            // for the tree, which seems to properly preserve type-use annotations
+          if (memberReferenceTree != null || lambdaExpressionTree != null) {
+            // For a method reference or lambda, we get generic type arguments from the javac's
+            // inferred type for the tree, which seems to properly preserve type-use annotations
             paramNullness =
                 GenericsChecks.getGenericMethodParameterNullness(
-                    i, overriddenMethod, ASTHelpers.getType(memberReferenceTree), state, config);
+                    i,
+                    overriddenMethod,
+                    ASTHelpers.getType(
+                        memberReferenceTree != null ? memberReferenceTree : lambdaExpressionTree),
+                    state,
+                    config);
           } else {
             // Use the enclosing class of the overriding method to find generic type arguments
             paramNullness =

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

@@ -820,16 +820,22 @@ public static Nullness getGenericMethodReturnTypeNullness(
   private static @Nullable Type getTypeForSymbol(Symbol symbol, VisitorState state, Config config) {
     if (symbol.isAnonymous()) {
       // For anonymous classes, symbol.type does not contain annotations on generic type parameters.
-      // So, we get a correct type from the enclosing NewClassTree.
+      // So, we get a correct type from the enclosing NewClassTree representing the anonymous class.
       TreePath path = state.getPath();
-      NewClassTree newClassTree = ASTHelpers.findEnclosingNode(path, NewClassTree.class);
-      if (newClassTree == null) {
+      path = ASTHelpers.findPathFromEnclosingNodeToTopLevel(path, NewClassTree.class);
+      NewClassTree newClassTree = (NewClassTree) path.getLeaf();
+      if (newClassTree == null || newClassTree.getClassBody() == null) {
         throw new RuntimeException(
-            "method should be inside a NewClassTree " + state.getSourceForNode(path.getLeaf()));
+            "method should be directly inside an anonymous NewClassTree "
+                + state.getSourceForNode(path.getLeaf()));
       }
       Type typeFromTree = getTreeType(newClassTree, config);
       if (typeFromTree != null) {
-        verify(state.getTypes().isAssignable(symbol.type, typeFromTree));
+        verify(
+            state.getTypes().isAssignable(symbol.type, typeFromTree),
+            "%s is not assignable to %s",
+            symbol.type,
+            typeFromTree);
       }
       return typeFromTree;
     } else {

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

@@ -507,6 +507,32 @@ public void testForMethodReferenceInAnAssignment() {
         .doTest();
   }
 
+  @Test
+  public void testForLambdaInAssignment() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "import java.util.function.Supplier;",
+            "class Test {",
+            "  interface A<T1 extends @Nullable Object> {",
+            "    String function(T1 o);",
+            "  }",
+            "  static String foo(Object o) {",
+            "    return o.toString();",
+            "  }",
+            "  static void testPositive() {",
+            "    // BUG: Diagnostic contains: dereferenced expression x is @Nullable",
+            "    A<@Nullable Object> p = x -> x.toString();",
+            "  }",
+            "  static void testNegative() {",
+            "    A<Object> p = x -> \"hello\";",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   @Test
   public void testForMethodReferenceForClassFieldAssignment() {
     makeHelper()
@@ -2298,6 +2324,35 @@ public void issue1129() {
         .doTest();
   }
 
+  @Test
+  public void issue1156() {
+    makeHelper()
+        .addSourceLines(
+            "Foo.java",
+            "import org.jspecify.annotations.NullMarked;",
+            "import java.util.function.Function;",
+            "import java.util.function.Supplier;",
+            "@NullMarked",
+            "public class Foo implements Supplier<Integer> {",
+            "  public Foo(Function<Integer, Integer> func) {",
+            "  }",
+            "  @Override",
+            "  public Integer get() {",
+            "    return 0;",
+            "  }",
+            "  public static void test() {",
+            "    new Supplier<Boolean>() {",
+            "      @Override",
+            "      public Boolean get() {",
+            "        Foo foo = new Foo(x -> 1);",
+            "        return true;",
+            "      }",
+            "    };",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   private CompilationTestHelper makeHelper() {
     return makeTestHelperWithArgs(
         Arrays.asList(