From c45bee532711449307ada603b293739a68f19b36 Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Wed, 4 May 2016 18:19:34 +0300
Subject: [PATCH] KT-12040 "Replace when with if" produce invalid code for
 first entry which has comment

 #KT-12040 Fixed
---
 .../kotlin/idea/util/CommentSaver.kt          | 21 +++++++++++++++++++
 .../branched/ifWhen/whenToIf/kt12040.kt       | 11 ++++++++++
 .../branched/ifWhen/whenToIf/kt12040.kt.after | 10 +++++++++
 .../intentions/IntentionTestGenerated.java    |  6 ++++++
 4 files changed, 48 insertions(+)
 create mode 100644 idea/testData/intentions/branched/ifWhen/whenToIf/kt12040.kt
 create mode 100644 idea/testData/intentions/branched/ifWhen/whenToIf/kt12040.kt.after

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/CommentSaver.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/CommentSaver.kt
index 76a75d36119..35cfed29943 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/CommentSaver.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/CommentSaver.kt
@@ -184,7 +184,12 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
         get() = getCopyableUserData(SAVED_TREE_KEY)
         set(value) = putCopyableUserData(SAVED_TREE_KEY, value)
 
+    var isFinished = false
+        private set
+
     fun deleteCommentsInside(element: PsiElement) {
+        assert(!isFinished)
+
         element.accept(object : PsiRecursiveElementVisitor() {
             override fun visitComment(comment: PsiComment) {
                 val treeElement = comment.savedTreeElement
@@ -196,6 +201,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
     }
 
     fun elementCreatedByText(createdElement: PsiElement, original: PsiElement, rangeInOriginal: TextRange) {
+        assert(!isFinished)
         assert(createdElement.textLength == rangeInOriginal.length)
         assert(createdElement.text == original.text.substring(rangeInOriginal.startOffset, rangeInOriginal.endOffset))
 
@@ -235,6 +241,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
     }
 
     fun restore(resultElements: PsiChildRange, forceAdjustIndent: Boolean = false) {
+        assert(!isFinished)
         assert(!resultElements.isEmpty)
 
         if (commentsToRestore.isNotEmpty() || lineBreaksToRestore.isNotEmpty()) {
@@ -281,6 +288,8 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
             }
             CodeStyleManager.getInstance(project).adjustLineIndent(file, resultElements.textRange)
         }
+
+        isFinished = true
     }
 
     private fun restoreComments(resultElements: PsiChildRange) {
@@ -302,6 +311,18 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks:
                     if (commentTreeElement.spaceBefore.isNotEmpty()) {
                         parent.addAfter(psiFactory.createWhiteSpace(commentTreeElement.spaceBefore), anchorElement)
                     }
+
+                    // make sure that there is a line break after EOL_COMMENT
+                    if (restored.tokenType == KtTokens.EOL_COMMENT) {
+                        val whiteSpace = restored.nextLeaf(skipEmptyElements = true) as? PsiWhiteSpace
+                        if (whiteSpace == null) {
+                            parent.addAfter(psiFactory.createWhiteSpace("\n"), restored)
+                        }
+                        else if (!whiteSpace.textContains('\n')) {
+                            val newWhiteSpace = psiFactory.createWhiteSpace("\n" + whiteSpace.text)
+                            whiteSpace.replace(newWhiteSpace)
+                        }
+                    }
                 }
                 else {
                     restored = parent.addBefore(comment, anchorElement) as PsiComment
diff --git a/idea/testData/intentions/branched/ifWhen/whenToIf/kt12040.kt b/idea/testData/intentions/branched/ifWhen/whenToIf/kt12040.kt
new file mode 100644
index 00000000000..27c57741df1
--- /dev/null
+++ b/idea/testData/intentions/branched/ifWhen/whenToIf/kt12040.kt
@@ -0,0 +1,11 @@
+fun foo(a: Int) {
+    <caret>when (a) {
+        // some comment
+        0 -> bar(a)
+
+        // another comment
+        1 -> bar(a)
+    }
+}
+
+fun bar(p: Int){}
\ No newline at end of file
diff --git a/idea/testData/intentions/branched/ifWhen/whenToIf/kt12040.kt.after b/idea/testData/intentions/branched/ifWhen/whenToIf/kt12040.kt.after
new file mode 100644
index 00000000000..21564a4448e
--- /dev/null
+++ b/idea/testData/intentions/branched/ifWhen/whenToIf/kt12040.kt.after
@@ -0,0 +1,10 @@
+fun foo(a: Int) {
+    if (a
+            // some comment
+            == 0) bar(a)
+
+    // another comment
+    else if (a == 1) bar(a)
+}
+
+fun bar(p: Int){}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index 46b5d53b0e1..2c374ac6211 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -1675,6 +1675,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     doTest(fileName);
                 }
 
+                @TestMetadata("kt12040.kt")
+                public void testKt12040() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/whenToIf/kt12040.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("whenWithDotQualifiedExpression.kt")
                 public void testWhenWithDotQualifiedExpression() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/whenToIf/whenWithDotQualifiedExpression.kt");