From d67941b84607868ec713da719d0a0943afce5f2b Mon Sep 17 00:00:00 2001
From: Suzanne Millstein <smillst@cs.washington.edu>
Date: Wed, 31 May 2023 08:49:58 -0700
Subject: [PATCH] Handle lambdas and method references in switch expressions

---
 .../framework/type/AnnotatedTypeFactory.java  | 13 +++++++++++--
 .../tests/all-systems/java17/Issue5930.java   | 19 +++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)
 create mode 100644 framework/tests/all-systems/java17/Issue5930.java

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 dd77d801623..a12249c4878 100644
--- a/framework/src/main/java/org/checkerframework/framework/type/AnnotatedTypeFactory.java
+++ b/framework/src/main/java/org/checkerframework/framework/type/AnnotatedTypeFactory.java
@@ -19,6 +19,7 @@
 import com.sun.source.tree.NewClassTree;
 import com.sun.source.tree.ReturnTree;
 import com.sun.source.tree.Tree;
+import com.sun.source.tree.Tree.Kind;
 import com.sun.source.tree.TypeCastTree;
 import com.sun.source.tree.VariableTree;
 import com.sun.source.util.TreePath;
@@ -4601,8 +4602,8 @@ public Pair<AnnotatedTypeMirror, AnnotatedExecutableType> getFnInterfaceFromTree
    * @return the functional interface type or an uninferred type argument
    */
   private AnnotatedTypeMirror getFunctionalInterfaceType(Tree tree) {
-
-    Tree parentTree = getPath(tree).getParentPath().getLeaf();
+    TreePath parentPath = getPath(tree).getParentPath();
+    Tree parentTree = parentPath.getLeaf();
     switch (parentTree.getKind()) {
       case PARENTHESIZED:
         return getFunctionalInterfaceType(parentTree);
@@ -4713,8 +4714,16 @@ private AnnotatedTypeMirror getFunctionalInterfaceType(Tree tree) {
             AnnotatedTypes.leastUpperBound(this, trueType, falseType);
         assertIsFunctionalInterface(conditionalType.getUnderlyingType(), parentTree, tree);
         return conditionalType;
+      case CASE:
+        // Get the functional interface type of the whole switch expression.
+        Tree switchTree = parentPath.getParentPath().getLeaf();
+        return getFunctionalInterfaceType(switchTree);
 
       default:
+        if (parentTree.getKind().toString().equals("YIELD")) {
+          TreePath pathToCase = TreePathUtil.pathTillOfKind(parentPath, Kind.CASE);
+          return getFunctionalInterfaceType(pathToCase.getParentPath().getLeaf());
+        }
         throw new BugInCF(
             "Could not find functional interface from assignment context. "
                 + "Unexpected tree type: "
diff --git a/framework/tests/all-systems/java17/Issue5930.java b/framework/tests/all-systems/java17/Issue5930.java
new file mode 100644
index 00000000000..0071aec2728
--- /dev/null
+++ b/framework/tests/all-systems/java17/Issue5930.java
@@ -0,0 +1,19 @@
+import java.util.function.Supplier;
+
+// @below-java17-jdk-skip-test
+public final class Issue5930 {
+  enum TestEnum {
+    FIRST,
+    SECOND
+  };
+
+  public static void main(String[] args) {
+    TestEnum testEnum = TestEnum.FIRST;
+    Supplier<Integer> supplier =
+        switch (testEnum) {
+          case FIRST -> () -> 1;
+          case SECOND -> () -> 2;
+        };
+    System.out.println(supplier.get());
+  }
+}