From ce54a2d4bd64a799f3bf03d4933f21fc971dd2d0 Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Thu, 11 Feb 2016 21:41:24 +0300
Subject: [PATCH] If to when conversion preserves comments too

---
 .../branchedTransformations/intentions/IfToWhenIntention.kt | 6 +++++-
 .../testData/intentions/branched/ifWhen/ifToWhen/comment.kt | 4 ++++
 .../intentions/branched/ifWhen/ifToWhen/comment.kt.after    | 6 ++++++
 .../kotlin/idea/intentions/IntentionTestGenerated.java      | 6 ++++++
 4 files changed, 21 insertions(+), 1 deletion(-)
 create mode 100644 idea/testData/intentions/branched/ifWhen/ifToWhen/comment.kt
 create mode 100644 idea/testData/intentions/branched/ifWhen/ifToWhen/comment.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/intentions/IfToWhenIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/intentions/IfToWhenIntention.kt
index 362334c4b58..6e2bf927ae8 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/intentions/IfToWhenIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/intentions/IfToWhenIntention.kt
@@ -21,6 +21,7 @@ import com.intellij.openapi.util.TextRange
 import org.jetbrains.kotlin.idea.intentions.SelfTargetingRangeIntention
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.getSubjectToIntroduce
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.introduceSubject
+import org.jetbrains.kotlin.idea.util.CommentSaver
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import java.util.*
@@ -32,6 +33,8 @@ class IfToWhenIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfExpres
     }
 
     override fun applyTo(element: KtIfExpression, editor: Editor?) {
+        val commentSaver = CommentSaver(element)
+
         var whenExpression = KtPsiFactory(element).buildExpression {
             appendFixedText("when {\n")
 
@@ -71,7 +74,8 @@ class IfToWhenIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfExpres
             whenExpression = whenExpression.introduceSubject()
         }
 
-        element.replace(whenExpression)
+        val result = element.replace(whenExpression)
+        commentSaver.restore(result)
     }
 
     private fun MutableList<KtExpression>.addOrBranches(expression: KtExpression): List<KtExpression> {
diff --git a/idea/testData/intentions/branched/ifWhen/ifToWhen/comment.kt b/idea/testData/intentions/branched/ifWhen/ifToWhen/comment.kt
new file mode 100644
index 00000000000..7819fde19da
--- /dev/null
+++ b/idea/testData/intentions/branched/ifWhen/ifToWhen/comment.kt
@@ -0,0 +1,4 @@
+fun foo(b: Boolean) {
+    <caret>if (b) 1 // 1
+    else 2
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/branched/ifWhen/ifToWhen/comment.kt.after b/idea/testData/intentions/branched/ifWhen/ifToWhen/comment.kt.after
new file mode 100644
index 00000000000..68d1fba3248
--- /dev/null
+++ b/idea/testData/intentions/branched/ifWhen/ifToWhen/comment.kt.after
@@ -0,0 +1,6 @@
+fun foo(b: Boolean) {
+    when {
+        b -> 1 // 1
+        else -> 2
+    }
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index 5455ed61d89..af6479ae2e6 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -1495,6 +1495,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/ifWhen/ifToWhen"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
                 }
 
+                @TestMetadata("comment.kt")
+                public void testComment() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/comment.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("ifWithEqualityTests.kt")
                 public void testIfWithEqualityTests() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/ifWithEqualityTests.kt");