Use CommentSaver instead of manual comment handling

idea/src/org/jetbrains/kotlin/idea/intentions/IfNullToElvisIntention.kt

@@ -18,15 +18,15 @@ package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
-import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiWhiteSpace
-import com.intellij.psi.PsiElement
-import com.intellij.psi.PsiRecursiveElementVisitor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.inspections.IntentionBasedInspection
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.expressionComparedToNull
+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.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
@@ -34,7 +34,6 @@ import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.typeUtil.isNothing
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
-import java.util.*
 
 class IfNullToElvisInspection : IntentionBasedInspection<KtIfExpression>(IfNullToElvisIntention())
 
@@ -63,13 +62,9 @@ class IfNullToElvisIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfE
             else -> null
         }
 
-        // do not loose any comments!
+        val childRangeBefore = PsiChildRange(declaration, element)
-        val comments = element.extractComments(ifNullExpr)
+        val commentSaver = CommentSaver(childRangeBefore)
-
+        val childRangeAfter = childRangeBefore.withoutLastStatement()
-        for (comment in comments) {
-            declaration.add(factory.createWhiteSpace())
-            declaration.add(comment)
-        }
 
         val elvis = factory.createExpressionByPattern("$0 ?: $1", initializer, ifNullExpr) as KtBinaryExpression
         val newElvis = initializer.replaced(elvis)
@@ -79,6 +74,8 @@ class IfNullToElvisIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfE
             declaration.setType(explicitTypeToSet)
         }
 
+        commentSaver.restore(childRangeAfter)
+
         editor?.caretModel?.moveToOffset(newElvis.right!!.textOffset)
     }
 
@@ -112,18 +109,8 @@ class IfNullToElvisIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfE
         }
     }
 
-    private fun PsiElement.extractComments(skipElement: PsiElement): List<PsiComment> {
+    private fun PsiChildRange.withoutLastStatement(): PsiChildRange {
-        val comments = ArrayList<PsiComment>()
+        val newLast = last!!.siblings(forward = false, withItself = false).first { it !is PsiWhiteSpace }
-        accept(object : PsiRecursiveElementVisitor() {
+        return PsiChildRange(first, newLast)
-            override fun visitElement(element: PsiElement) {
-                if (element == skipElement) return
-                super.visitElement(element)
-            }
-
-            override fun visitComment(comment: PsiComment) {
-                comments.add(comment)
-            }
-        })
-        return comments
     }
 }

idea/testData/intentions/ifNullToElvis/CommentInBlock.kt.after

@@ -1,4 +1,5 @@
 fun foo(p: List<String?>): Int {
-    val v = p[0] ?: <caret>return -1 // return -1 if null
+    val v = p[0] ?: // return -1 if null
+            <caret>return -1
     return v.length
 }

idea/testData/intentions/ifNullToElvis/Comments.kt.after

@@ -1,5 +1,5 @@
 fun foo(p: List<String?>): Int {
-    val v = p[0] ?: <caret>return -1 /* null */
+    val v = p[0] ?: /* null */<caret>return -1
     // now check if v is null
     // return -1
     return v.length

idea/testData/intentions/ifNullToElvis/Comments2.kt.after

@@ -1,7 +1,7 @@
 fun foo(p: List<String?>): Int {
-    val v = p[0] ?: return -1 // v is null
+    val v = p[0] ?: // v is null
             // we should do something with it
-            // let's return -1
+            return -1 // let's return -1
     // end of if
     return v.length
 }