From d97f659b90cc77672330da3997b423fe84db3c59 Mon Sep 17 00:00:00 2001
From: Mikhail Pyltsin <Mikhail.Pyltsin@jetbrains.com>
Date: Thu, 14 Dec 2023 15:56:04 +0100
Subject: [PATCH] [java-inspection] IDEA-340970 don't offer fix if `yield`
 statement is broken

GitOrigin-RevId: 414d7230cd3887b5a92664582a42c8775ccf753a
---
 .../EnhancedSwitchMigrationInspection.java      |  7 +++++--
 .../beforeBrokenSwitchYieldStatement.java       | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)
 create mode 100644 java/java-tests/testData/inspection/switchExpressionMigration/beforeBrokenSwitchYieldStatement.java

diff --git a/java/java-impl/src/com/intellij/codeInspection/EnhancedSwitchMigrationInspection.java b/java/java-impl/src/com/intellij/codeInspection/EnhancedSwitchMigrationInspection.java
index b176eec83a71..6f1f92c74483 100644
--- a/java/java-impl/src/com/intellij/codeInspection/EnhancedSwitchMigrationInspection.java
+++ b/java/java-impl/src/com/intellij/codeInspection/EnhancedSwitchMigrationInspection.java
@@ -574,7 +574,7 @@ public final class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLoc
       PsiStatement statement = result[i];
       if (statement instanceof PsiReturnStatement returnStatement) {
         PsiExpression returnValue = returnStatement.getReturnValue();
-        if (returnValue == null) {
+        if (returnValue == null || PsiTreeUtil.hasErrorElements(returnValue)) {
           return null;
         }
         result[i] = createYieldStatement(returnValue);
@@ -585,7 +585,7 @@ public final class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLoc
       Collection<PsiReturnStatement> returnStatements = PsiTreeUtil.findChildrenOfType(copy, PsiReturnStatement.class);
       for (PsiReturnStatement returnStatement : returnStatements) {
         PsiExpression returnValue = returnStatement.getReturnValue();
-        if (returnValue == null) {
+        if (returnValue == null || PsiTreeUtil.hasErrorElements(returnValue)) {
           return null;
         }
         returnStatement.replace(createYieldStatement(returnValue));
@@ -763,6 +763,9 @@ public final class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLoc
           result = new SwitchRuleExpressionResult(rExpression);
         }
         else {
+          if (PsiTreeUtil.hasErrorElements(rExpression)) {
+            return null;
+          }
           result = new SwitchStatementBranch(withLastStatementReplacedWithYield(statements, rExpression));
         }
       }
diff --git a/java/java-tests/testData/inspection/switchExpressionMigration/beforeBrokenSwitchYieldStatement.java b/java/java-tests/testData/inspection/switchExpressionMigration/beforeBrokenSwitchYieldStatement.java
new file mode 100644
index 000000000000..32983cb60cc5
--- /dev/null
+++ b/java/java-tests/testData/inspection/switchExpressionMigration/beforeBrokenSwitchYieldStatement.java
@@ -0,0 +1,17 @@
+// "Replace with 'switch' expression" "false"
+import java.util.*;
+
+class Switch {
+  void test(Integer o) {
+    String s = "";
+    s<caret>witch (o) {
+      case 1:
+        s = "";
+        break;
+      default:
+        System.out.println("1");
+        s = {System.out.println("break");
+        break;
+        } ;
+    }
+  }
\ No newline at end of file