IDEA-203984 "Remove unreachable branches" quickfix breaks code on switch expression

java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DataFlowInspectionBase.java

@@ -314,21 +314,7 @@ public class DataFlowInspectionBase extends AbstractBaseJavaLocalInspectionTool
     for (PsiExpression label : trueLabels) {
       PsiSwitchLabelStatementBase labelStatement = Objects.requireNonNull(PsiImplUtil.getSwitchLabel(label));
       PsiSwitchBlock statement = labelStatement.getEnclosingSwitchBlock();
-      if (statement == null) continue;
-      if (PsiTreeUtil.getChildrenOfTypeAsList(statement.getBody(), PsiSwitchLabelStatementBase.class).size() == 1 &&
-          Objects.requireNonNull(labelStatement.getCaseValues()).getExpressionCount() == 1) {
-        boolean canUnwrap;
-        if (statement instanceof PsiSwitchStatement) {
-          canUnwrap = BreakConverter.from(statement) != null;
-        }
-        else {
-          canUnwrap = labelStatement instanceof PsiSwitchLabeledRuleStatement &&
-                      ((PsiSwitchLabeledRuleStatement)labelStatement).getBody() instanceof PsiExpressionStatement;
-        }
-        if (!canUnwrap) {
-          continue;
-        }
-      }
+      if (statement == null || !canRemoveUnreachableBranches(labelStatement, statement)) continue;
       if (!StreamEx.iterate(labelStatement, Objects::nonNull, l -> PsiTreeUtil.getPrevSiblingOfType(l, PsiSwitchLabelStatementBase.class))
         .skip(1).map(PsiSwitchLabelStatementBase::getCaseValues)
         .nonNull().flatArray(PsiExpressionList::getExpressions).allMatch(falseLabels::contains)) {
@@ -347,6 +333,21 @@ public class DataFlowInspectionBase extends AbstractBaseJavaLocalInspectionTool
     }
   }
 
+  private static boolean canRemoveUnreachableBranches(PsiSwitchLabelStatementBase labelStatement, PsiSwitchBlock statement) {
+    if (Objects.requireNonNull(labelStatement.getCaseValues()).getExpressionCount() != 1) return true;
+    List<PsiSwitchLabelStatementBase> allBranches =
+      PsiTreeUtil.getChildrenOfTypeAsList(statement.getBody(), PsiSwitchLabelStatementBase.class);
+    if (statement instanceof PsiSwitchStatement) {
+      // Cannot do anything if we have already single branch and we cannot restore flow due to non-terminal breaks
+      return allBranches.size() != 1 || BreakConverter.from(statement) != null;
+    }
+    // Expression switch: if we cannot unwrap existing branch and the other one is default case, we cannot kill it either 
+    return (allBranches.size() <= 2 &&
+           !allBranches.stream().allMatch(branch -> branch == labelStatement || branch.isDefaultCase())) ||
+           (labelStatement instanceof PsiSwitchLabeledRuleStatement &&
+            ((PsiSwitchLabeledRuleStatement)labelStatement).getBody() instanceof PsiExpressionStatement);
+  }
+
   private void reportConstants(ProblemReporter reporter, DataFlowInstructionVisitor visitor) {
     visitor.getConstantExpressions().forEach((expression, result) -> {
       if (result == ConstantResult.UNKNOWN) return;

java/java-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/UnwrapSwitchLabelFix.java

@@ -36,10 +36,12 @@ public class UnwrapSwitchLabelFix implements LocalQuickFix {
     PsiSwitchBlock block = labelStatement.getEnclosingSwitchBlock();
     if (block == null) return;
     List<PsiSwitchLabelStatementBase> labels = PsiTreeUtil.getChildrenOfTypeAsList(block.getBody(), PsiSwitchLabelStatementBase.class);
+    boolean shouldKeepDefault = block instanceof PsiSwitchExpression &&
+                                !(labelStatement instanceof PsiSwitchLabeledRuleStatement &&
+                                  ((PsiSwitchLabeledRuleStatement)labelStatement).getBody() instanceof PsiExpressionStatement);
     for (PsiSwitchLabelStatementBase otherLabel : labels) {
-      if (otherLabel != labelStatement) {
-        DeleteSwitchLabelFix.deleteLabel(otherLabel);
-      }
+      if (otherLabel == labelStatement || (shouldKeepDefault && otherLabel.isDefaultCase())) continue;
+      DeleteSwitchLabelFix.deleteLabel(otherLabel);
     }
     for (PsiExpression expression : Objects.requireNonNull(labelStatement.getCaseValues()).getExpressions()) {
       if (expression != label) {

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/deleteSwitchLabel/afterUnwrapExpressionJava12Preview.java

@@ -13,7 +13,7 @@ class Main {
                 }
                 //6
                 //7
-                //other
+                default -> "and more"; //other
             });
         }
     }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/deleteSwitchLabel/beforeUnwrapExpression3Java12Preview.java

@@ -0,0 +1,14 @@
+// "Remove unreachable branches" "false"
+class X {
+    int x(E e) {
+        // We cannot do anything here (yet)
+        return switch (E.AA) {
+            case A<caret>A: 
+                System.out.println(9);
+            default: break 0;
+        };
+    }
+}
+enum E {
+    AA,BB
+}