Introduce/eliminate when subject intentions to keep line breaks and comments

idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/branchedTransformationUtils.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.intentions.branchedTransformations
 
 import org.jetbrains.kotlin.idea.core.replaced
+import org.jetbrains.kotlin.idea.util.CommentSaver
 import org.jetbrains.kotlin.idea.util.psi.patternMatching.matches
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
@@ -100,6 +101,8 @@ private fun KtExpression?.getWhenConditionSubjectCandidate(): KtExpression? {
 fun KtWhenExpression.introduceSubject(): KtWhenExpression {
     val subject = getSubjectToIntroduce()!!
 
+    val commentSaver = CommentSaver(this, saveLineBreaks = true)
+
     val whenExpression = KtPsiFactory(this).buildExpression {
         appendFixedText("when(").appendExpression(subject).appendFixedText("){\n")
 
@@ -126,7 +129,9 @@ fun KtWhenExpression.introduceSubject(): KtWhenExpression {
         appendFixedText("}")
     } as KtWhenExpression
 
-    return replaced(whenExpression)
+    val result = replaced(whenExpression)
+    commentSaver.restore(result)
+    return result
 }
 
 private fun BuilderByPattern<KtExpression>.appendConditionWithSubjectRemoved(conditionExpression: KtExpression?, subject: KtExpression) {

idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/intentions/EliminateWhenSubjectIntention.kt

@@ -20,6 +20,7 @@ import com.intellij.codeInsight.intention.LowPriorityAction
 import com.intellij.openapi.editor.Editor
 import org.jetbrains.kotlin.idea.intentions.SelfTargetingIntention
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.toExpression
+import org.jetbrains.kotlin.idea.util.CommentSaver
 import org.jetbrains.kotlin.psi.KtNameReferenceExpression
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.KtWhenExpression
@@ -36,6 +37,8 @@ class EliminateWhenSubjectIntention : SelfTargetingIntention<KtWhenExpression>(K
     override fun applyTo(element: KtWhenExpression, editor: Editor?) {
         val subject = element.subjectExpression!!
 
+        val commentSaver = CommentSaver(element, saveLineBreaks = true)
+
         val whenExpression = KtPsiFactory(element).buildExpression {
             appendFixedText("when {\n")
 
@@ -57,6 +60,7 @@ class EliminateWhenSubjectIntention : SelfTargetingIntention<KtWhenExpression>(K
             appendFixedText("}")
         }
 
-        element.replace(whenExpression)
+        val result = element.replace(whenExpression)
+        commentSaver.restore(result)
     }
 }

idea/testData/intentions/branched/when/eliminateSubject/lineBreaksAndComments.kt

@@ -0,0 +1,13 @@
+class Klass<T>
+
+fun test(obj: Any): String {
+    return <caret>when (obj) {
+        is String -> "string" // return "string"
+
+        // it's an Int
+        is Int -> "int"
+
+        // otherwise
+        else -> "unknown"
+    }
+}

idea/testData/intentions/branched/when/eliminateSubject/lineBreaksAndComments.kt.after

@@ -0,0 +1,13 @@
+class Klass<T>
+
+fun test(obj: Any): String {
+    return when {
+        obj is String -> "string" // return "string"
+
+    // it's an Int
+        obj is Int -> "int"
+
+    // otherwise
+        else -> "unknown"
+    }
+}

idea/testData/intentions/branched/when/introduceSubject/lineBreaksAndComments.kt

@@ -0,0 +1,13 @@
+class Klass<T>
+
+fun test(obj: Any): String {
+    return <caret>when {
+        obj is String -> "string" // return "string"
+
+        // it's an Int
+        obj is Int -> "int"
+
+        // otherwise
+        else -> "unknown"
+    }
+}

idea/testData/intentions/branched/when/introduceSubject/lineBreaksAndComments.kt.after

@@ -0,0 +1,13 @@
+class Klass<T>
+
+fun test(obj: Any): String {
+    return when (obj) {
+        is String -> "string" // return "string"
+
+    // it's an Int
+        is Int -> "int"
+
+    // otherwise
+        else -> "unknown"
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -2074,6 +2074,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/when/eliminateSubject"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
                 }
 
+                @TestMetadata("lineBreaksAndComments.kt")
+                public void testLineBreaksAndComments() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/when/eliminateSubject/lineBreaksAndComments.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("whenWithEqualityTests.kt")
                 public void testWhenWithEqualityTests() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/when/eliminateSubject/whenWithEqualityTests.kt");
@@ -2158,6 +2164,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/when/introduceSubject"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
                 }
 
+                @TestMetadata("lineBreaksAndComments.kt")
+                public void testLineBreaksAndComments() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/when/introduceSubject/lineBreaksAndComments.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("whenWithEqualityTests.kt")
                 public void testWhenWithEqualityTests() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/when/introduceSubject/whenWithEqualityTests.kt");