[java-inspections] Support yield in push down (IDEA-301920)

GitOrigin-RevId: a6d2e3c37656ae2998665d1128f2585e0a91ccb6

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/controlflow/SwitchExpressionCanBePushedDownInspection.java

@@ -11,12 +11,15 @@ import com.intellij.openapi.project.Project;
 import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.ArrayUtil;
+import com.intellij.util.containers.ContainerUtil;
 import com.siyeh.InspectionGadgetsBundle;
+import com.siyeh.ig.psiutils.ControlFlowUtils;
 import com.siyeh.ig.psiutils.EquivalenceChecker;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.Objects;
 
@@ -89,8 +92,23 @@ public class SwitchExpressionCanBePushedDownInspection extends AbstractBaseJavaL
       if (!(statement instanceof PsiSwitchLabeledRuleStatement rule)) return null;
       PsiStatement ruleBody = rule.getBody();
       if (ruleBody instanceof PsiThrowStatement) continue;
-      if (!(ruleBody instanceof PsiExpressionStatement expr)) return null;
-      result.add(expr.getExpression());
+      if (ruleBody instanceof PsiBlockStatement blockStatement) {
+        if (!(block instanceof PsiSwitchExpression)) return null;
+        if (ControlFlowUtils.codeBlockMayCompleteNormally(blockStatement.getCodeBlock())) return null;
+        Collection<PsiYieldStatement> yields = PsiTreeUtil.findChildrenOfType(blockStatement, PsiYieldStatement.class);
+        List<PsiYieldStatement> myYields = ContainerUtil.filter(yields, st -> st.findEnclosingExpression() == block);
+        for (PsiYieldStatement yield : myYields) {
+          PsiExpression expression = yield.getExpression();
+          if (expression == null) return null;
+          result.add(expression);
+        }
+        continue;
+      }
+      if (ruleBody instanceof PsiExpressionStatement expr) {
+        result.add(expr.getExpression());
+        continue;
+      }
+      return null;
     }
     if (result.size() < 2) return null;
     return result;

plugins/InspectionGadgets/test/com/siyeh/igfixes/controlflow/push_down_switch/afterYield.java

@@ -0,0 +1,15 @@
+// "Push down 'switch' expression" "true-preview"
+class X {
+  String print(int value) {
+    return "[" + switch (value) {
+        case 0 -> "zero";
+        case 1 -> "one";
+        default -> {
+            if (value < 0) {
+                yield "negative";
+            }
+            yield value;
+        }
+    } + "]";
+  }
+}

plugins/InspectionGadgets/test/com/siyeh/igfixes/controlflow/push_down_switch/beforeYield.java

@@ -0,0 +1,15 @@
+// "Push down 'switch' expression" "true-preview"
+class X {
+  String print(int value) {
+    return <caret>switch (value) {
+      case 0 -> "[" + "zero" + "]";
+      case 1 -> "[" + "one" + "]";
+      default -> {
+        if (value < 0) {
+          yield "[" + "negative" + "]";
+        }
+        yield "[" + value + "]";
+      }
+    };
+  }
+}

plugins/InspectionGadgets/testsrc/com/siyeh/ig/fixes/controlflow/SwitchExpressionCanBePushedDownInspectionTest.java

@@ -4,11 +4,14 @@ package com.siyeh.ig.fixes.controlflow;
 import com.intellij.codeInsight.daemon.quickFix.LightQuickFixParameterizedTestCase;
 import com.intellij.codeInspection.LocalInspectionTool;
 import com.intellij.openapi.application.ex.PathManagerEx;
+import com.intellij.testFramework.LightProjectDescriptor;
 import com.siyeh.ig.LightJavaInspectionTestCase;
 import com.siyeh.ig.controlflow.SwitchExpressionCanBePushedDownInspection;
 import com.siyeh.ig.controlflow.SwitchStatementsWithoutDefaultInspection;
 import org.jetbrains.annotations.NotNull;
 
+import static com.intellij.testFramework.fixtures.LightJavaCodeInsightFixtureTestCase.JAVA_17;
+
 public class SwitchExpressionCanBePushedDownInspectionTest extends LightQuickFixParameterizedTestCase {
   @Override
   protected LocalInspectionTool @NotNull [] configureLocalInspectionTools() {