Add PureExceptLambda annotation (#1325)

### Description
This PR adds support for nullability preserving methods annotated with a
new `PureExceptLambda` annotation.
When a method is marked as `@PureExceptLambda`, it means that there's no
way that it could effect the variables used within a lambda passed to it
(no side-effects except those possibly performed by the lambda once
invoked), and that the lambda will only be called within that method and
not stored for later use. Allowing developers to implement their own
lambda api's without needing explicit checks in NullAway.

Here's an example of a simple adapter pattern with the issue:
```java
@FunctionalInterface
public interface LegacyListener {
  void onEvent(String payload);
}

public static class EventAdapter {
  @PureExceptLambda
  public void withLegacyListener(String payload, LegacyListener listener) {
    listener.onEvent(payload);
  }
}

public class Service {
  private @nullable Database db;  // set elsewhere

  public void process(String input) {
    if (db == null) {
      // After this, `db` is known non-null in this branch
      throw new IllegalStateException("db not initialized");
    }
    EventAdapter.withLegacyListener(input, payload -> db.execute(payload));
  }
}
```
Before this PR, NullAway would have warned that `db` is null inside the
lambda returned from `EventSource.withListener`. However, now its seen
as nullness preserving, it preserves the non-null state.

<details>
<summary>Details</summary>

Please note that once you click "Create Pull Request" you will be asked
to sign our [Uber Contributor License
Agreement](https://cla-assistant.io/uber/NullAway) via [CLA
assistant](https://cla-assistant.io/).

Before pressing the "Create Pull Request" button, please provide the
following:

- [x] A description about what and why you are contributing, even if
it's trivial.

- [x] The issue number(s) or PR number(s) in the description if you are
contributing in response to those.

  - [x] If applicable, unit tests.

</details>

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added @PureExceptLambda annotation to mark methods that preserve
nullness properties when invoking lambda parameters, enabling improved
null-safety analysis for lambda-based APIs.

* **Tests**
* Added test coverage validating nullness fact preservation in methods
annotated with @PureExceptLambda.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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

Author: FxMorin

annotations/src/main/java/com/uber/nullaway/annotations/PureExceptLambda.java

@@ -0,0 +1,52 @@
+package com.uber.nullaway.annotations;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Indicates that a method or constructor preserves the nullness environment when invoking lambda
+ * expressions passed as parameters.
+ *
+ * <p>You can only use this on methods which have a single lambda callback. Otherwise, the purity
+ * cannot be guaranteed.
+ *
+ * <p>When a method or constructor is annotated with {@code @PureExceptLambda}, NullAway assumes
+ * that any lambda arguments passed to it are invoked <em>synchronously</em> and that the invocation
+ * does not modify the visible state or capture the lambda for later (asynchronous) execution. As a
+ * result, NullAway treats the body of the lambda as being executed in the same nullness environment
+ * as the code at the call site.
+ *
+ * <p>In other words, this annotation tells NullAway that:
+ *
+ * <ul>
+ *   <li>The annotated method or constructor does not perform side effects on fields or global state
+ *       visible at the call site.
+ *   <li>The single functional parameter (e.g., lambda or method reference) is invoked synchronously
+ *       within the method or constructor body and is not stored or invoked later.
+ *   <li>Nullability information from the surrounding context is preserved inside the lambda body.
+ * </ul>
+ *
+ * <p>This allows NullAway to safely reason about nullness within lambdas passed to such methods as
+ * if the lambda were executed inline at the call site.
+ *
+ * <p><b>Important:</b> NullAway does <em>not</em> verify that annotated methods actually satisfy
+ * these requirements. The annotation represents a contract that developers must uphold manually.
+ * Misuse of this annotation (e.g., annotating methods that perform side effects or invoke lambdas
+ * asynchronously) can lead to unsound nullability analysis. Note that this annotation requires only
+ * side-effect freedom; determinism (always producing the same result for the same inputs) is not
+ * required.
+ *
+ * <p><b>Example:</b>
+ *
+ * <pre>{@code
+ * @PureExceptLambda
+ * public static <T> void runWith(T value, Consumer<T> action) {
+ *     action.accept(value);
+ * }
+ * }</pre>
+ */
+@Retention(RetentionPolicy.CLASS)
+@Target({ElementType.METHOD, ElementType.CONSTRUCTOR})
+public @interface PureExceptLambda {}

nullaway/src/main/java/com/uber/nullaway/handlers/SynchronousCallbackHandler.java

@@ -19,6 +19,8 @@
 import com.uber.nullaway.LibraryModels.MethodRef;
 import com.uber.nullaway.dataflow.AccessPath;
 import java.util.function.Predicate;
+import javax.lang.model.element.AnnotationMirror;
+import org.checkerframework.nullaway.javacutil.AnnotationUtils;
 
 public class SynchronousCallbackHandler extends BaseNoOpHandler {
 
@@ -67,6 +69,11 @@ public Predicate<AccessPath> getAccessPathPredicateForNestedMethod(
         // preserve access paths for all callbacks passed to stream methods
         return TRUE_AP_PREDICATE;
       }
+      // If the callee is a method that preserves the nullability of lambdas passed as parameters.
+      // (e.g., annotated with @PureExceptLambda).
+      if (isMethodPureExceptLambda(symbol)) {
+        return TRUE_AP_PREDICATE;
+      }
       String invokedMethodName = symbol.getSimpleName().toString();
       if (METHOD_NAME_TO_SIG_AND_PARAM_INDEX.containsKey(invokedMethodName)) {
         ImmutableMap<MethodRef, Integer> entriesForMethodName =
@@ -91,4 +98,16 @@ public Predicate<AccessPath> getAccessPathPredicateForNestedMethod(
     }
     return FALSE_AP_PREDICATE;
   }
+
+  private boolean isMethodPureExceptLambda(Symbol.MethodSymbol methodSymbol) {
+    for (AnnotationMirror annotation : methodSymbol.getAnnotationMirrors()) {
+      String name = AnnotationUtils.annotationName(annotation);
+      int lastDot = name.lastIndexOf('.');
+      String simpleName = lastDot == -1 ? name : name.substring(lastDot + 1);
+      if (simpleName.equals("PureExceptLambda")) {
+        return true;
+      }
+    }
+    return false;
+  }
 }

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

@@ -0,0 +1,84 @@
+package com.uber.nullaway;
+
+import java.util.Arrays;
+import org.junit.Test;
+
+/**
+ * Tests for preserving environment nullness facts for lambdas passed to nullness preserving methods
+ * (methods annotated with @PureExceptLambda). Variables captured in such lambdas should not be
+ * treated as automatically nullable; instead, the facts at the call site should be visible in the
+ * lambda body.
+ */
+public class PureExceptLambdaTests extends NullAwayTestsBase {
+
+  @Test
+  public void methodLambdaPreservesFacts() {
+    makeTestHelperWithArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber"))
+        .addSourceLines(
+            "com/example/library/PureLibrary.java",
+            "package com.example.library;",
+            "import com.uber.nullaway.annotations.PureExceptLambda;",
+            "import java.util.function.Consumer;",
+            "public class PureLibrary {",
+            "  @PureExceptLambda",
+            "  public static <T> void withConsumer(T t, Consumer<T> consumer) {",
+            "    consumer.accept(t);",
+            "  }",
+            "}")
+        .addSourceLines(
+            "com/uber/Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "import com.example.library.PureLibrary;",
+            "public class Test {",
+            "  private @Nullable Object f;",
+            "  public void test1() {",
+            "    if (this.f == null) {",
+            "      throw new IllegalArgumentException();",
+            "    }",
+            "    // f is known non-null after the check; that fact should be visible in the lambda",
+            "    PureLibrary.withConsumer(\"x\", v -> System.out.println(this.f.toString()));",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void methodLambdaDoesNotPreserveFacts() {
+    makeTestHelperWithArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber"))
+        .addSourceLines(
+            "com/example/library/OrdinaryLibrary.java",
+            "package com.example.library;",
+            "import java.util.function.Consumer;",
+            "public class OrdinaryLibrary {",
+            "  public static <T> void withConsumer(T t, Consumer<T> consumer) {",
+            "    consumer.accept(t);",
+            "  }",
+            "}")
+        .addSourceLines(
+            "com/uber/Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "import com.example.library.OrdinaryLibrary;",
+            "public class Test {",
+            "  private @Nullable Object f;",
+            "  public void test1() {",
+            "    if (this.f == null) {",
+            "      throw new IllegalArgumentException();",
+            "    }",
+            "    // Without nullness preserving annotation, we should not propagate the non-null fact into the lambda",
+            "    // BUG: Diagnostic contains: dereferenced expression this.f is @Nullable",
+            "    OrdinaryLibrary.withConsumer(\"x\", v -> System.out.println(this.f.toString()));",
+            "  }",
+            "}")
+        .doTest();
+  }
+}