EnhancedSwitchMigration: generate proper assignment, skip default branch if exists: IDEA-208192

java/java-impl/src/com/intellij/codeInspection/EnhancedSwitchMigrationInspection.java

@@ -451,11 +451,10 @@ public class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLocalInsp
                                                               boolean isExhaustive) {
     PsiElement parent = statement.getParent();
     PsiElement anchor = parent instanceof PsiLabeledStatement ? parent : statement;
-    PsiDeclarationStatement declaration = tryCast(PsiTreeUtil.getPrevSiblingOfType(anchor, PsiStatement.class), PsiDeclarationStatement.class);
+    PsiVariable assignedVariable = null;
     List<SwitchExpressionBranch> newBranches = new ArrayList<>();
-    PsiLocalVariable assignedVariable = getVariable(declaration);
-    boolean isRightAfterDeclaration = assignedVariable != null;
     boolean hasAssignedBranch = false;
+    boolean wasDefault = false;
     for (OldSwitchStatementBranch branch : branches) {
       if (!isConvertibleBranch(branch, false)) return null;
       if (branch.isFallthrough() && branch.getStatements().length == 0) continue;
@@ -482,11 +481,15 @@ public class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLocalInsp
         hasAssignedBranch = true;
         result = new SwitchRuleExpressionResult(rExpression);
       }
-      newBranches.add(new SwitchExpressionBranch(branch.isDefault(), branch.getCaseExpressions(), result, branch.getRelatedStatements()));
+      boolean isDefault = branch.isDefault();
+      if (isDefault) {
+        wasDefault = true;
+      }
+      newBranches.add(new SwitchExpressionBranch(isDefault, branch.getCaseExpressions(), result, branch.getRelatedStatements()));
     }
     if (assignedVariable == null || !hasAssignedBranch) return null;
     PsiExpression initializer = assignedVariable.getInitializer();
-    if (initializer != null) {
+    if (initializer != null && !wasDefault) {
       newBranches.add(new SwitchExpressionBranch(true,
                                                  Collections.emptyList(),
                                                  new SwitchRuleExpressionResult(initializer),
@@ -494,9 +497,24 @@ public class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLocalInsp
     } else {
       if (!isExhaustive) return null;
     }
+    boolean isRightAfterDeclaration = isRightAfterDeclaration(anchor, assignedVariable);
     return new SwitchExistingVariableReplacer(assignedVariable, statement, expressionBeingSwitched, newBranches, isRightAfterDeclaration);
   }
 
+  private static boolean isRightAfterDeclaration(PsiElement anchor, PsiVariable assignedVariable) {
+    PsiDeclarationStatement declaration = tryCast(PsiTreeUtil.getPrevSiblingOfType(anchor, PsiStatement.class), PsiDeclarationStatement.class);
+    if (declaration != null) {
+      PsiElement[] elements = declaration.getDeclaredElements();
+      if (elements.length == 1) {
+        PsiLocalVariable localVariable = tryCast(elements[0], PsiLocalVariable.class);
+        if (localVariable != null && localVariable == assignedVariable) {
+          return true;
+        }
+      }
+    }
+    return false;
+  }
+
   /**
    * Replaces with enhanced switch statement
    */

java/java-tests/testData/inspection/switchExpressionMigration/afterSwitchVarAssignmentMultiVar.java

@@ -0,0 +1,16 @@
+// "Replace with 'switch' expression" "true"
+import java.util.*;
+
+class SwitchExpressionMigration {
+  private static void m() {
+    String s = "foo";
+    int x = 5;
+    System.out.println(s);
+
+      s = switch (x) {
+          case 1 -> "bar";
+          case 2 -> "baz";
+          default -> "zuq";
+      };
+  }
+}

java/java-tests/testData/inspection/switchExpressionMigration/beforeSwitchVarAssignmentMultiVar.java

@@ -0,0 +1,16 @@
+// "Replace with 'switch' expression" "true"
+import java.util.*;
+
+class SwitchExpressionMigration {
+  private static void m() {
+    String s = "foo";
+    int x = 5;
+    System.out.println(s);
+
+    switch<caret> (x) {
+      case 1: s = "bar";break;
+      case 2: s = "baz";break;
+      default: s = "zuq";break;
+    }
+  }
+}