Handle lambdas and method references in switch expressions

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: "

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());
+  }
+}