Move variable declaration into when: fix comment place

#KT-31362 Fixed

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

idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment2.kt

@@ -0,0 +1,7 @@
+fun foo(style: Int?): Int {
+    val a<caret> = style ?: 0 // comment
+    return when (a) {
+        0 -> 0
+        else -> a
+    }
+}

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
+    }
+}

idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment3.kt

@@ -0,0 +1,6 @@
+fun foo(style: Int?): Int {
+    val a<caret> = style ?: 0 // comment
+    return when (a) { 0 -> 0
+        else -> a
+    }
+}

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
+    }
+}

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