KT-12040 "Replace when with if" produce invalid code for first entry which has comment

#KT-12040 Fixed

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

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){}

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){}

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");