Handle lambdas and method references in switch expressions

Author: smillst

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