From 1703f867dfefee753ce94208187b80b2d728d5c3 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <Tagir.Valeev@jetbrains.com>
Date: Wed, 6 Feb 2019 11:48:19 +0700
Subject: [PATCH] Fix IDEA-206725 "Insert default branch" produces uncompilable
 code

---
 .../com/siyeh/ig/fixes/CreateDefaultBranchFix.java   | 12 +++++++++++-
 .../controlflow/create_default/afterAssignment.java  | 10 ++++++++++
 .../controlflow/create_default/beforeAssignment.java |  8 ++++++++
 3 files changed, 29 insertions(+), 1 deletion(-)
 create mode 100644 plugins/InspectionGadgets/test/com/siyeh/igfixes/controlflow/create_default/afterAssignment.java
 create mode 100644 plugins/InspectionGadgets/test/com/siyeh/igfixes/controlflow/create_default/beforeAssignment.java

diff --git a/plugins/InspectionGadgets/src/com/siyeh/ig/fixes/CreateDefaultBranchFix.java b/plugins/InspectionGadgets/src/com/siyeh/ig/fixes/CreateDefaultBranchFix.java
index 574a277fba11..fa91306b9ecd 100644
--- a/plugins/InspectionGadgets/src/com/siyeh/ig/fixes/CreateDefaultBranchFix.java
+++ b/plugins/InspectionGadgets/src/com/siyeh/ig/fixes/CreateDefaultBranchFix.java
@@ -10,6 +10,7 @@ import com.intellij.ide.fileTemplates.JavaTemplateUtil;
 import com.intellij.openapi.editor.Editor;
 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.IncorrectOperationException;
 import com.siyeh.ig.psiutils.ControlFlowUtils;
@@ -22,6 +23,7 @@ import java.io.IOException;
 import java.util.*;
 
 public class CreateDefaultBranchFix extends BaseSwitchFix {
+  private static final String PLACEHOLDER_NAME = "$EXPRESSION$";
   private final String myMessage;
 
   public CreateDefaultBranchFix(@NotNull PsiSwitchBlock block, String message) {
@@ -86,7 +88,7 @@ public class CreateDefaultBranchFix extends BaseSwitchFix {
     FileTemplate branchTemplate = FileTemplateManager.getInstance(project).getCodeTemplate(JavaTemplateUtil.TEMPLATE_SWITCH_DEFAULT_BRANCH);
     Properties props = FileTemplateManager.getInstance(project).getDefaultProperties();
     PsiExpression expression = switchBlock.getExpression();
-    props.setProperty(FileTemplate.ATTRIBUTE_EXPRESSION, expression == null ? "" : expression.getText());
+    props.setProperty(FileTemplate.ATTRIBUTE_EXPRESSION, PLACEHOLDER_NAME);
     PsiType expressionType = expression == null ? null : expression.getType();
     props.setProperty(FileTemplate.ATTRIBUTE_EXPRESSION_TYPE, expressionType == null ? "" : expressionType.getCanonicalText());
     PsiStatement statement;
@@ -99,6 +101,14 @@ public class CreateDefaultBranchFix extends BaseSwitchFix {
         }
       }
       statement = JavaPsiFacade.getElementFactory(project).createStatementFromText("{" + text + "}", switchBlock);
+      if (expression != null) {
+        PsiElement[] refs = PsiTreeUtil.collectElements(
+          statement, e -> e instanceof PsiReferenceExpression && e.textMatches(PLACEHOLDER_NAME));
+        for (PsiElement ref : refs) {
+          // This would add parentheses when necessary
+          ref.replace(expression);
+        }
+      }
     }
     catch (IOException | IncorrectOperationException e) {
       throw new IncorrectOperationException("Incorrect file template", (Throwable)e);
diff --git a/plugins/InspectionGadgets/test/com/siyeh/igfixes/controlflow/create_default/afterAssignment.java b/plugins/InspectionGadgets/test/com/siyeh/igfixes/controlflow/create_default/afterAssignment.java
new file mode 100644
index 000000000000..87d6c2c43a2b
--- /dev/null
+++ b/plugins/InspectionGadgets/test/com/siyeh/igfixes/controlflow/create_default/afterAssignment.java
@@ -0,0 +1,10 @@
+// "Insert 'default' branch" "true"
+class X {
+  void test(int i, int j) {
+    switch(i=j) {
+      case 0:break;
+        default:
+            throw new IllegalStateException("Unexpected value: " + (i = j));
+    }
+  }
+}
\ No newline at end of file
diff --git a/plugins/InspectionGadgets/test/com/siyeh/igfixes/controlflow/create_default/beforeAssignment.java b/plugins/InspectionGadgets/test/com/siyeh/igfixes/controlflow/create_default/beforeAssignment.java
new file mode 100644
index 000000000000..7cbc25adb29e
--- /dev/null
+++ b/plugins/InspectionGadgets/test/com/siyeh/igfixes/controlflow/create_default/beforeAssignment.java
@@ -0,0 +1,8 @@
+// "Insert 'default' branch" "true"
+class X {
+  void test(int i, int j) {
+    switch(i=j) {
+      case 0:<caret>break;
+    }
+  }
+}
\ No newline at end of file