diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/IfNullToElvisIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/IfNullToElvisIntention.kt index 060f3a875f1..57133dae32b 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/IfNullToElvisIntention.kt +++ b/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.PsiElement -import com.intellij.psi.PsiRecursiveElementVisitor +import com.intellij.psi.PsiWhiteSpace 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(IfNullToElvisIntention()) @@ -63,13 +62,9 @@ class IfNullToElvisIntention : SelfTargetingRangeIntention(KtIfE else -> null } - // do not loose any comments! - val comments = element.extractComments(ifNullExpr) - - for (comment in comments) { - declaration.add(factory.createWhiteSpace()) - declaration.add(comment) - } + val childRangeBefore = PsiChildRange(declaration, element) + val commentSaver = CommentSaver(childRangeBefore) + val childRangeAfter = childRangeBefore.withoutLastStatement() val elvis = factory.createExpressionByPattern("$0 ?: $1", initializer, ifNullExpr) as KtBinaryExpression val newElvis = initializer.replaced(elvis) @@ -79,6 +74,8 @@ class IfNullToElvisIntention : SelfTargetingRangeIntention(KtIfE declaration.setType(explicitTypeToSet) } + commentSaver.restore(childRangeAfter) + editor?.caretModel?.moveToOffset(newElvis.right!!.textOffset) } @@ -112,18 +109,8 @@ class IfNullToElvisIntention : SelfTargetingRangeIntention(KtIfE } } - private fun PsiElement.extractComments(skipElement: PsiElement): List { - val comments = ArrayList() - accept(object : PsiRecursiveElementVisitor() { - override fun visitElement(element: PsiElement) { - if (element == skipElement) return - super.visitElement(element) - } - - override fun visitComment(comment: PsiComment) { - comments.add(comment) - } - }) - return comments + private fun PsiChildRange.withoutLastStatement(): PsiChildRange { + val newLast = last!!.siblings(forward = false, withItself = false).first { it !is PsiWhiteSpace } + return PsiChildRange(first, newLast) } } \ No newline at end of file diff --git a/idea/testData/intentions/ifNullToElvis/CommentInBlock.kt.after b/idea/testData/intentions/ifNullToElvis/CommentInBlock.kt.after index 37d306d64bd..ac0a8592d10 100644 --- a/idea/testData/intentions/ifNullToElvis/CommentInBlock.kt.after +++ b/idea/testData/intentions/ifNullToElvis/CommentInBlock.kt.after @@ -1,4 +1,5 @@ fun foo(p: List): Int { - val v = p[0] ?: return -1 // return -1 if null + val v = p[0] ?: // return -1 if null + return -1 return v.length } \ No newline at end of file diff --git a/idea/testData/intentions/ifNullToElvis/Comments.kt.after b/idea/testData/intentions/ifNullToElvis/Comments.kt.after index 45251facf01..edb1c6a992b 100644 --- a/idea/testData/intentions/ifNullToElvis/Comments.kt.after +++ b/idea/testData/intentions/ifNullToElvis/Comments.kt.after @@ -1,5 +1,5 @@ fun foo(p: List): Int { - val v = p[0] ?: return -1 /* null */ + val v = p[0] ?: /* null */return -1 // now check if v is null // return -1 return v.length diff --git a/idea/testData/intentions/ifNullToElvis/Comments2.kt.after b/idea/testData/intentions/ifNullToElvis/Comments2.kt.after index 173a507b95c..876c50fc712 100644 --- a/idea/testData/intentions/ifNullToElvis/Comments2.kt.after +++ b/idea/testData/intentions/ifNullToElvis/Comments2.kt.after @@ -1,7 +1,7 @@ fun foo(p: List): 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 } \ No newline at end of file