switch expressions 2.0: update CFG & inference

GitOrigin-RevId: 8f83e7c2c57819e34c4cb8fa5d44a863b0b58532

java/java-psi-api/src/com/intellij/psi/LambdaUtil.java

@@ -123,6 +123,12 @@ public class LambdaUtil {
         return isValidLambdaContext(element.getParent());
       }
     }
+    if (context instanceof PsiYieldStatement) {
+      PsiSwitchExpression switchExpression = ((PsiYieldStatement)context).findEnclosingExpression();
+      if (switchExpression != null) {
+        return isValidLambdaContext(switchExpression.getParent());
+      }
+    }
     if (context instanceof PsiExpressionStatement) {
       PsiElement parent = context.getParent();
       if (parent instanceof PsiSwitchLabeledRuleStatement) {

java/java-psi-api/src/com/intellij/psi/impl/source/resolve/graphInference/PsiPolyExpressionUtil.java

@@ -86,6 +86,7 @@ public class PsiPolyExpressionUtil {
     PsiElement parent = PsiUtil.skipParenthesizedExprUp(expr).getParent();
     if (parent instanceof PsiExpressionStatement && parent.getParent() instanceof PsiSwitchLabeledRuleStatement || 
         parent instanceof PsiBreakStatement || 
+        parent instanceof PsiYieldStatement ||
         parent instanceof PsiThrowStatement) {
       PsiSwitchExpression switchExpression = PsiTreeUtil.getParentOfType(expr, PsiSwitchExpression.class, true, PsiMember.class, PsiLambdaExpression.class);
       return switchExpression  != null &&

java/java-psi-api/src/com/intellij/psi/util/PsiUtil.java

@@ -382,7 +382,7 @@ public final class PsiUtil extends PsiUtilCore {
       ContainerUtil.addIfNotNull(result, statement.getExpression());
     }
     List<PsiYieldStatement> yields = new ArrayList<>();
-    addStatements(yields, container, PsiYieldStatement.class, element -> false);
+    addStatements(yields, container, PsiYieldStatement.class, element -> element instanceof PsiSwitchExpression);
     for (PsiYieldStatement statement : yields) {
       ContainerUtil.addIfNotNull(result, statement.getExpression());
     }

java/java-psi-impl/src/com/intellij/psi/controlFlow/ControlFlowAnalyzer.java

@@ -396,10 +396,18 @@ class ControlFlowAnalyzer extends JavaElementVisitor {
 
   @Override
   public void visitBreakStatement(PsiBreakStatement statement) {
+    generateYieldInstructions(statement, statement.getValueExpression(), statement.findExitedElement());
+  }
+
+  @Override
+  public void visitYieldStatement(PsiYieldStatement statement) {
+    generateYieldInstructions(statement, statement.getExpression(), statement.findEnclosingExpression());
+  }
+
+  private void generateYieldInstructions(PsiStatement statement, PsiExpression valueExpression, PsiElement exitedStatement) {
     startElement(statement);
-    generateExpressionInstructions(statement.getValueExpression());
+    generateExpressionInstructions(valueExpression);
 
-    PsiElement exitedStatement = statement.findExitedElement();
     if (exitedStatement != null) {
       callFinallyBlocksOnExit(exitedStatement);
 

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/switchExpressions/EnhancedSwitchDefinitelyAssigned.java

@@ -118,8 +118,8 @@ class C {
 
   static class FinalFieldValueBreakSwitchExpression {
     final String s = switch ((int)Math.random()) {
-      case 1: break "a";
-      default: break "?";
+      case 1: yield "a";
+      default: yield "?";
     };
     {
       System.out.println(s);
@@ -136,8 +136,8 @@ class C {
 
   void finalVariableValueBreakSwitchExpression(String s) {
     final int n = switch (s) {
-      case "a": break 1;
-      default: break 0;
+      case "a": yield 1;
+      default: yield 0;
     };
     System.out.println(n);
   }
@@ -163,8 +163,8 @@ class C {
   void definitelyAssignedInSwitchExpressionValueBreak(String s) {
     int n;
     int x = switch (s) {
-      case "a": break n = 1;
-      default: break n = 0;
+      case "a": yield n = 1;
+      default: yield n = 0;
     };
     System.out.println(n);
   }
@@ -172,8 +172,8 @@ class C {
   void notDefinitelyAssignedInSwitchExpressionValueBreak(String s) {
     int n;
     int x = switch (s) {
-      case "a": break n = 1;
-      default: break 0;
+      case "a": yield n = 1;
+      default: yield 0;
     };
     System.out.println(<error descr="Variable 'n' might not have been initialized">n</error>);
   }
@@ -198,7 +198,7 @@ class C {
       String string = switch (n) {
         case -1 -> throw new RuntimeException();
         case 0 -> s = "a";
-        default -> { break s = "b"; }
+        default -> { yield s = "b"; }
       };
     }
     System.out.println(s);

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/switchExpressions/EnhancedSwitchUnreachable.java

@@ -47,7 +47,7 @@ class C {
 
   void endlessLoopInBranchWithValueBreak(String arg) {
     int result = switch (arg) {
-      case "one" -> { while(true); <error descr="Unreachable statement">break 1;</error>}
+      case "one" -> { while(true); <error descr="Unreachable statement">yield 1;</error>}
       default -> 0;
     };
     System.out.println(result);
@@ -73,8 +73,8 @@ class C {
     int n;
     return 2;
     n = <error descr="Unreachable statement">switch</error>(s) {
-      case "a": n = 1; break 1;
-      default: n = 0; break 0;
+      case "a": n = 1; yield 1;
+      default: n = 0; yield 0;
     };
   }
 
@@ -129,8 +129,8 @@ class C {
     int foo(String s) throws Exception {
       try {
         return switch (s) {
-          case "a": break bar(1);
-          default: break bar(0);
+          case "a": yield bar(1);
+          default: yield bar(0);
         };
       } finally {
         System.out.println("b");

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/switchExpressions/SimpleInferenceCases.java

@@ -23,31 +23,31 @@ no instance(s) of type variable(s) exist so that Object conforms to String">foo(
             default -> "str";
         });
     String s3 = foo(() -> switch (i) {default -> bar();});
-    String s4 = foo(() -> switch (i) {default -> { break bar();}});
+    String s4 = foo(() -> switch (i) {default -> { yield bar();}});
     String s5 = foo(<error descr="Incompatible types. Required String but 'foo' was inferred to T:
-no instance(s) of type variable(s) exist so that Integer conforms to String">() -> switch (i) {default -> { break 1;}}</error>);
+no instance(s) of type variable(s) exist so that Integer conforms to String">() -> switch (i) {default -> { yield 1;}}</error>);
     String s6 = switch (i) {
       case 1 -> <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">2</error>;
       default -> {
-        break <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
+        yield <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
       }
     };
     Supplier<String> stringSupplier = switch (i) {
             default -> {
-                break () -> <error descr="Bad return type in lambda expression: int cannot be converted to String">1</error>;
+                yield () -> <error descr="Bad return type in lambda expression: int cannot be converted to String">1</error>;
             }
         };
     String s7 = switch (i) {
       case 1: {
-        break <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
+        yield <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
       }
       default: {
         int i1 = switch (0) {
           default -> {
-            break 1;
+            yield 1;
           }
         };
-        break "";
+        yield "";
       }
     };
   }
@@ -56,15 +56,15 @@ no instance(s) of type variable(s) exist so that Integer conforms to String">()
     String s = switch (i) {
       default -> switch (0) {
         default -> {
-          break <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
+          yield <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
         }
       };
     };
     String s1 = switch (i) {
       default -> {
-        break switch (0) {
+        yield switch (0) {
           default -> {
-            break <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
+            yield <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
           }
         };
       }
@@ -72,9 +72,9 @@ no instance(s) of type variable(s) exist so that Integer conforms to String">()
 
     String s2 = switch (i) {
       default: {
-        break switch (0) {
+        yield switch (0) {
           default -> {
-            break <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
+            yield <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
           }
         };
       }
@@ -82,24 +82,24 @@ no instance(s) of type variable(s) exist so that Integer conforms to String">()
 
     String s3 = switch (0) {
       default: {
-        break switch (1) {
+        yield switch (1) {
           case 2: {
             System.out.println();
             int inside_switch = switch (8) {
               default:
-                break 1;
+                yield 1;
             };
           }
           case 1: 
-            if (i > 3) break <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">3</error>;
+            if (i > 3) yield <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">3</error>;
           case 0:
             try {
-              break <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">42</error>;
+              yield <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">42</error>;
             } finally {
               //do nothing
             }
           default:
-            break <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
+            yield <error descr="Bad type in switch expression: int cannot be converted to java.lang.String">1</error>;
         };
       }
     };

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/switchExpressions/SwitchExpressionHasResult.java

@@ -14,7 +14,7 @@ class C {
   void defaultBranchSometimesHasNoResult(int n, boolean b) {
     String s = <error descr="Switch expression should produce result in all execution paths">switch</error> (n) {
       default: {
-        if (b) break "";
+        if (b) yield "";
       }
     };
   }
@@ -22,7 +22,7 @@ class C {
   void defaultRuleSometimesHasNoResult(int n, boolean b) {
     String s = switch (n) {
       <error descr="Switch expression rule should produce result in all execution paths">default</error> -> {
-        if (b) break "";
+        if (b) yield "";
       }
     };
   }
@@ -43,7 +43,7 @@ class C {
     String s = switch (n) {
       default:
         if (b) throw new RuntimeException();
-        break "";
+        yield "";
     };
   }
 
@@ -51,7 +51,7 @@ class C {
     String s = switch (n) {
       default -> {
         if (b) throw new RuntimeException();
-        break "";
+        yield "";
       }
     };
   }
@@ -59,9 +59,9 @@ class C {
   void defaultBranchHasManyResults(int n, int k) {
     String s = switch (n) {
       default: {
-        if (k < n) break "a";
-        if (k > n) break "b";
-        break "c";
+        if (k < n) yield "a";
+        if (k > n) yield "b";
+        yield "c";
       }
     };
   }
@@ -69,16 +69,16 @@ class C {
   void defaultRuleHasManyResults(int n, int k) {
     String s = switch (n) {
       default -> {
-        if (k < n) break "a";
-        if (k > n) break "b";
-        break "c";
+        if (k < n) yield "a";
+        if (k > n) yield "b";
+        yield "c";
       }
     };
   }
 
   void oneOfBranchesHasNoResult(int n) {
     String s = <error descr="Switch expression should produce result in all execution paths">switch</error> (n) {
-      case 0: break "";
+      case 0: yield "";
       default:
     };
   }
@@ -110,8 +110,8 @@ class C {
   void allBranchesDoHaveResult(int n) {
     String s = switch (n) {
       case -1: throw new RuntimeException();
-      case 0: break "a";
-      default: break "b";
+      case 0: yield "a";
+      default: yield "b";
     };
   }
 
@@ -129,8 +129,8 @@ class C {
     } finally {
       s = switch (n) {
         case -1: throw new RuntimeException();
-        case 0: break "a";
-        default: break "b";
+        case 0: yield "a";
+        default: yield "b";
       };
     }
   }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/switchExpressions/SwitchExpressions.java

@@ -12,7 +12,7 @@ class SwitchExpressions {
 
     System.out.println(switch (new Random().nextInt()) {
       case 0 -> throw new IllegalStateException("no args");
-      <error descr="Different case kinds used in the switch">case 1:</error> break "lone";
+      <error descr="Different case kinds used in the switch">case 1:</error> yield "lone";
     });
 
     System.out.println(

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/switchExpressions/ValueBreaks.java

@@ -9,19 +9,19 @@ class ValueBreaks {
     <error descr="Value break outside switch expression">break 42;</error>
 
     switch (0) {
-      case 0: <error descr="Value break outside switch expression">break 42;</error>
+      case 0: <error descr="Yield outside of switch expression">yield 42;</error>
       case 1: break <error descr="Undefined label: 'ref'">ref</error>;
       case 2: break <error descr="Undefined label: 'wtf'">wtf</error>;
       case 3: ref: break ref;
-      case 4: ref: <error descr="Value break outside switch expression">break (ref);</error>
+      case 4: ref: <error descr="Yield outside of switch expression">yield (ref);</error>
     };
 
     sink(switch (0) {
       case 0 -> { while (true) <error descr="Value break outside switch expression">break 42;</error> }
       case 1 -> { while (true) <error descr="Value break outside switch expression">break ref;</error> }
       case 2 -> { while (true) break <error descr="Undefined label: 'wtf'">wtf</error>; }
-      case 3 -> { break ref; }
-      case 4 -> { break (ref); }
+      case 3 -> { yield ref; }
+      case 4 -> { yield (ref); }
       case 5 -> { break <error descr="Cannot resolve symbol 'wtf'">wtf</error>; }
       case 6 -> {
         int a = 0;
@@ -48,7 +48,7 @@ class ValueBreaks {
 
     while (true) {
       sink(switch (0) {
-        default: <error descr="Missing break value">break;</error>
+        default: yield<error descr="Expression expected">;</error>
       });
     }
   }

java/java-tests/testSrc/com/intellij/java/codeInsight/daemon/JavaSwitchExpressionsHighlightingTest.kt

@@ -2,13 +2,10 @@
 package com.intellij.java.codeInsight.daemon
 
 import com.intellij.JavaTestUtil
-import com.intellij.openapi.roots.LanguageLevelModuleExtension
-import com.intellij.openapi.roots.ModuleRootModificationUtil
-import com.intellij.pom.java.LanguageLevel
 import com.intellij.testFramework.fixtures.LightJavaCodeInsightFixtureTestCase
 
 class JavaSwitchExpressionsHighlightingTest : LightJavaCodeInsightFixtureTestCase() {
-  override fun getProjectDescriptor() = JAVA_12
+  override fun getProjectDescriptor() = JAVA_13
   override fun getBasePath() = "${JavaTestUtil.getRelativeJavaTestDataPath()}/codeInsight/daemonCodeAnalyzer/switchExpressions"
 
   fun testEnhancedSwitchStatements() = doTest()
@@ -21,16 +18,7 @@ class JavaSwitchExpressionsHighlightingTest : LightJavaCodeInsightFixtureTestCas
   fun testEnhancedSwitchUnreachable() = doTest()
   fun testSwitchExpressionHasResult() = doTest()
 
-  fun testYieldStatements() = try {
-    level(LanguageLevel.JDK_13_PREVIEW)
-    doTest()
-  }
-  finally {
-    level(LanguageLevel.JDK_12_PREVIEW)
-  }
-
-  private fun level(level: LanguageLevel) =
-    ModuleRootModificationUtil.updateModel(module) { it.getModuleExtension(LanguageLevelModuleExtension::class.java).languageLevel = level }
+  fun testYieldStatements() = doTest()
 
   private fun doTest() {
     myFixture.configureByFile("${getTestName(false)}.java")