Fix library models check for expression-body lambdas

Author: dyrpsf

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

@@ -1165,9 +1165,18 @@ public Description matchLambdaExpression(LambdaExpressionTree tree, VisitorState
     updateEnvironmentMapping(state.getPath(), state);
     handler.onMatchLambdaExpression(
         tree, new MethodAnalysisContext(this, state, funcInterfaceMethod));
-    if (codeAnnotationInfo.isSymbolUnannotated(funcInterfaceMethod, config, handler)) {
-      return Description.NO_MATCH;
+    
+    boolean isUnannotated = codeAnnotationInfo.isSymbolUnannotated(funcInterfaceMethod, config, handler);
+    if (isUnannotated) {
+      // Even if the functional interface is unannotated, a library model may override its
+      // return type to be @NonNull. If so, we should still check the lambda expression.
+      Nullness returnNullness =
+          getMethodReturnNullness(funcInterfaceMethod, state, Nullness.NULLABLE);
+      if (!returnNullness.equals(Nullness.NONNULL)) {
+        return Description.NO_MATCH;
+      }
     }
+
     Description description =
         checkParamOverriding(
             tree.getParameters().stream().map(ASTHelpers::getSymbol).collect(Collectors.toList()),

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

@@ -0,0 +1,22 @@
+package com.uber.nullaway;
+
+import org.junit.Test;
+
+public class NullAwayLambdaTests extends NullAwayTestsBase {
+
+  @Test
+  public void testLambdaExpressionLibraryModelOverride() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.google.common.base.Function;",
+            "public class Test {",
+            "  public void test() {",
+            "    // BUG: Diagnostic contains: returning @Nullable expression from method with @NonNull return type",
+            "    Function<Object, Object> f1 = (x) -> null;",
+            "  }",
+            "}")
+        .doTest();
+  }
+}