diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/MoveVariableDeclarationIntoWhenInspection.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/MoveVariableDeclarationIntoWhenInspection.kt index 4de2c890c0a..e2abe8fbbad 100644 --- a/idea/src/org/jetbrains/kotlin/idea/inspections/MoveVariableDeclarationIntoWhenInspection.kt +++ b/idea/src/org/jetbrains/kotlin/idea/inspections/MoveVariableDeclarationIntoWhenInspection.kt @@ -11,13 +11,17 @@ import com.intellij.codeInspection.ProblemDescriptor import com.intellij.codeInspection.ProblemsHolder import com.intellij.openapi.project.Project import com.intellij.openapi.util.TextRange +import com.intellij.psi.PsiComment import com.intellij.psi.PsiElementVisitor +import com.intellij.psi.PsiWhiteSpace import com.intellij.psi.SmartPsiElementPointer import org.jetbrains.kotlin.idea.intentions.branchedTransformations.isOneLiner import org.jetbrains.kotlin.idea.intentions.loopToCallChain.countUsages import org.jetbrains.kotlin.idea.intentions.loopToCallChain.previousStatement +import org.jetbrains.kotlin.lexer.KtTokens import org.jetbrains.kotlin.psi.* import org.jetbrains.kotlin.psi.psiUtil.createSmartPointer +import org.jetbrains.kotlin.psi.psiUtil.siblings class MoveVariableDeclarationIntoWhenInspection : AbstractKotlinInspection(), CleanupLocalInspectionTool { override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean): PsiElementVisitor = @@ -96,6 +100,21 @@ private class VariableDeclarationIntoWhenFix( val property = descriptor.psiElement as? KtProperty ?: return val subjectExpression = subjectExpressionPointer.element ?: return val newElement = transform(property)?.copy() ?: return + + val lastChild = newElement.lastChild + if (lastChild is PsiComment && lastChild.node.elementType == KtTokens.EOL_COMMENT) { + val leftBrace = subjectExpression.siblings(withItself = false).firstOrNull { it.node.elementType == KtTokens.LBRACE } + val whiteSpaceBeforeComment = lastChild.prevSibling?.takeIf { it is PsiWhiteSpace } + if (leftBrace != null) { + subjectExpression.parent.addAfter(lastChild, leftBrace) + if (whiteSpaceBeforeComment != null) { + subjectExpression.parent.addAfter(whiteSpaceBeforeComment, leftBrace) + } + } + whiteSpaceBeforeComment?.delete() + lastChild.delete() + } + subjectExpression.replace(newElement) property.delete() } diff --git a/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment2.kt b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment2.kt new file mode 100644 index 00000000000..29682e19ff6 --- /dev/null +++ b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment2.kt @@ -0,0 +1,7 @@ +fun foo(style: Int?): Int { + val a = style ?: 0 // comment + return when (a) { + 0 -> 0 + else -> a + } +} \ No newline at end of file diff --git a/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment2.kt.after b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment2.kt.after new file mode 100644 index 00000000000..54bbe40c9ae --- /dev/null +++ b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment2.kt.after @@ -0,0 +1,6 @@ +fun foo(style: Int?): Int { + return when (val a = style ?: 0) { // comment + 0 -> 0 + else -> a + } +} \ No newline at end of file diff --git a/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment3.kt b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment3.kt new file mode 100644 index 00000000000..3dd0c797ff8 --- /dev/null +++ b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment3.kt @@ -0,0 +1,6 @@ +fun foo(style: Int?): Int { + val a = style ?: 0 // comment + return when (a) { 0 -> 0 + else -> a + } +} \ No newline at end of file diff --git a/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment3.kt.after b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment3.kt.after new file mode 100644 index 00000000000..54bbe40c9ae --- /dev/null +++ b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment3.kt.after @@ -0,0 +1,6 @@ +fun foo(style: Int?): Int { + return when (val a = style ?: 0) { // comment + 0 -> 0 + else -> a + } +} \ No newline at end of file diff --git a/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java index 939fe8dbdab..d4fde38c4db 100644 --- a/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java +++ b/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java @@ -4698,6 +4698,16 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest { runTest("idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment.kt"); } + @TestMetadata("withComment2.kt") + public void testWithComment2() throws Exception { + runTest("idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment2.kt"); + } + + @TestMetadata("withComment3.kt") + public void testWithComment3() throws Exception { + runTest("idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment3.kt"); + } + @TestMetadata("withNewLine.kt") public void testWithNewLine() throws Exception { runTest("idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withNewLine.kt");