From 4fcedc1441a451a77e8c08eac29af0110cb6f0c4 Mon Sep 17 00:00:00 2001
From: Michael Ernst <mernst@cs.washington.edu>
Date: Tue, 12 Dec 2023 10:43:02 -0800
Subject: [PATCH] Permit lambda parameters to have type `Optional`

---
 .../checker/optional/OptionalVisitor.java           |  9 ++++++++-
 checker/tests/optional/IfPresentRefinement.java     |  4 +---
 checker/tests/optional/OptionalParameterTest.java   | 13 +++++++++++++
 .../framework/type/AnnotatedTypeFactory.java        |  2 ++
 4 files changed, 24 insertions(+), 4 deletions(-)
 create mode 100644 checker/tests/optional/OptionalParameterTest.java

diff --git a/checker/src/main/java/org/checkerframework/checker/optional/OptionalVisitor.java b/checker/src/main/java/org/checkerframework/checker/optional/OptionalVisitor.java
index 143403d93d3..d24fa56b140 100644
--- a/checker/src/main/java/org/checkerframework/checker/optional/OptionalVisitor.java
+++ b/checker/src/main/java/org/checkerframework/checker/optional/OptionalVisitor.java
@@ -13,6 +13,7 @@
 import com.sun.source.tree.Tree.Kind;
 import com.sun.source.tree.UnaryTree;
 import com.sun.source.tree.VariableTree;
+import com.sun.source.util.TreePath;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
@@ -536,7 +537,13 @@ public Void visitVariable(VariableTree tree, Void p) {
       if (ekind.isField()) {
         checker.reportWarning(tree, "optional.field");
       } else if (ekind == ElementKind.PARAMETER) {
-        checker.reportWarning(tree, "optional.parameter");
+        TreePath paramPath = getCurrentPath();
+        Tree parent = paramPath.getParentPath().getLeaf();
+        if (parent.getKind() == Tree.Kind.LAMBDA_EXPRESSION) {
+          // Exception to rule: lambda parameters can have type Optional.
+        } else {
+          checker.reportWarning(tree, "optional.parameter");
+        }
       }
     }
     return super.visitVariable(tree, p);
diff --git a/checker/tests/optional/IfPresentRefinement.java b/checker/tests/optional/IfPresentRefinement.java
index ec78ec914d9..1b828bddb15 100644
--- a/checker/tests/optional/IfPresentRefinement.java
+++ b/checker/tests/optional/IfPresentRefinement.java
@@ -1,18 +1,16 @@
 import java.util.Optional;
 
+@SuppressWarnings("optional.parameter")
 public class IfPresentRefinement {
 
-  @SuppressWarnings("optional.parameter")
   void m1(Optional<String> o) {
     o.ifPresent(s -> o.get());
   }
 
-  @SuppressWarnings("optional.parameter")
   void m2(Optional<String> o) {
     o.ifPresentOrElse(s -> o.get(), () -> {});
   }
 
-  @SuppressWarnings("optional.parameter")
   void m3(Optional<String> o) {
     // :: error: (method.invocation)
     o.ifPresentOrElse(s -> o.get(), () -> o.get());
diff --git a/checker/tests/optional/OptionalParameterTest.java b/checker/tests/optional/OptionalParameterTest.java
new file mode 100644
index 00000000000..eb68d1d24a0
--- /dev/null
+++ b/checker/tests/optional/OptionalParameterTest.java
@@ -0,0 +1,13 @@
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+class OptionalParameterTest {
+  public void findDatesByIds2(List<Integer> ids) {
+    ids.stream()
+        .map(Optional::ofNullable)
+        .flatMap(optional -> optional.map(Stream::of).orElseGet(Stream::empty))
+        .collect(Collectors.toList());
+  }
+}
diff --git a/framework/src/main/java/org/checkerframework/framework/type/AnnotatedTypeFactory.java b/framework/src/main/java/org/checkerframework/framework/type/AnnotatedTypeFactory.java
index 5ce3a87c468..3475b8def08 100644
--- a/framework/src/main/java/org/checkerframework/framework/type/AnnotatedTypeFactory.java
+++ b/framework/src/main/java/org/checkerframework/framework/type/AnnotatedTypeFactory.java
@@ -3730,6 +3730,8 @@ public void setVisitorTreePath(@Nullable TreePath visitorTreePath) {
    * <p>Note that the given Tree has to be within the current compilation unit, otherwise null will
    * be returned.
    *
+   * <p>Within a subclass of BaseTypeVisitor, use {@code getCurrentPath()} rather than this method.
+   *
    * @param tree the {@link Tree} to get the path for
    * @return the path for {@code tree} under the current root. Returns null if {@code tree} is not
    *     within the current compilation unit.