From 8082a5daf7e40156a13c456c48cac226ad4fb0f7 Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Tue, 6 Mar 2018 04:52:45 +0300
Subject: [PATCH] Fix "Add braces to if" when semicolon is used instead of a
 new line

So #KT-23123 Fixed
---
 .../idea/intentions/AddBracesIntention.kt     | 19 ++++++++++++-------
 ...ddBracesForIfWithSemicolonAndExpression.kt |  5 +++++
 ...esForIfWithSemicolonAndExpression.kt.after |  8 ++++++++
 .../intentions/IntentionTestGenerated.java    |  6 ++++++
 4 files changed, 31 insertions(+), 7 deletions(-)
 create mode 100644 idea/testData/intentions/addBraces/addBracesForIfWithSemicolonAndExpression.kt
 create mode 100644 idea/testData/intentions/addBraces/addBracesForIfWithSemicolonAndExpression.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt
index 520248cadf4..7f1e278954c 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt
@@ -19,12 +19,11 @@ package org.jetbrains.kotlin.idea.intentions
 import com.intellij.openapi.editor.Editor
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.idea.refactoring.getLineNumber
 import org.jetbrains.kotlin.idea.util.CommentSaver
+import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
-import org.jetbrains.kotlin.psi.psiUtil.allChildren
-import org.jetbrains.kotlin.psi.psiUtil.getNextSiblingIgnoringWhitespace
-import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.psi.psiUtil.*
 import java.lang.IllegalArgumentException
 
 class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.java, "Add braces") {
@@ -53,8 +52,15 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
         if (editor == null) throw IllegalArgumentException("This intention requires an editor")
         val expression = element.getTargetExpression(editor.caretModel.offset)!!
 
-        if (element.nextSibling?.text == ";") {
-            element.nextSibling!!.delete()
+        val psiFactory = KtPsiFactory(element)
+
+        val semicolon = element.getNextSiblingIgnoringWhitespaceAndComments()?.takeIf { it.node.elementType == KtTokens.SEMICOLON }
+        if (semicolon != null) {
+            val afterSemicolon = semicolon.getNextSiblingIgnoringWhitespace()
+            if (semicolon.getLineNumber() == afterSemicolon?.getLineNumber())
+                semicolon.replace(psiFactory.createNewLine())
+            else
+                semicolon.delete()
         }
 
         val nextComment = when {
@@ -68,7 +74,6 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
         }
         nextComment?.delete()
 
-        val psiFactory = KtPsiFactory(element)
         val result = expression.replace(psiFactory.createSingleStatementBlock(expression))
 
         if (element is KtDoWhileExpression) { // remove new line between '}' and while
diff --git a/idea/testData/intentions/addBraces/addBracesForIfWithSemicolonAndExpression.kt b/idea/testData/intentions/addBraces/addBracesForIfWithSemicolonAndExpression.kt
new file mode 100644
index 00000000000..086d528d2a4
--- /dev/null
+++ b/idea/testData/intentions/addBraces/addBracesForIfWithSemicolonAndExpression.kt
@@ -0,0 +1,5 @@
+fun <T> doSomething(a: T) {}
+
+fun foo() {
+    <caret>if (true) doSomething("test"); doSomething("")
+}
diff --git a/idea/testData/intentions/addBraces/addBracesForIfWithSemicolonAndExpression.kt.after b/idea/testData/intentions/addBraces/addBracesForIfWithSemicolonAndExpression.kt.after
new file mode 100644
index 00000000000..9cd2e4b2dec
--- /dev/null
+++ b/idea/testData/intentions/addBraces/addBracesForIfWithSemicolonAndExpression.kt.after
@@ -0,0 +1,8 @@
+fun <T> doSomething(a: T) {}
+
+fun foo() {
+    if (true) {
+        doSomething("test")
+    }
+    doSomething("")
+}
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index 3c28a663150..8d8150da1a7 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -65,6 +65,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("addBracesForIfWithSemicolonAndExpression.kt")
+        public void testAddBracesForIfWithSemicolonAndExpression() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addBraces/addBracesForIfWithSemicolonAndExpression.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("addBracesForWhile.kt")
         public void testAddBracesForWhile() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addBraces/addBracesForWhile.kt");