From 5d0b42ee71176ad54ed429fbb4be0bbfcf621416 Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Tue, 26 May 2015 00:09:32 +0300
Subject: [PATCH] ConvertToForEachFunctionCallIntention preserves comments
 outside of body too

---
 .../intentions/ConvertToForEachFunctionCallIntention.kt  | 9 ++++++---
 .../convertToForEachFunctionCall/commentsOutsideBody.kt  | 3 +++
 .../commentsOutsideBody.kt.after                         | 3 +++
 .../kotlin/idea/intentions/IntentionTestGenerated.java   | 6 ++++++
 4 files changed, 18 insertions(+), 3 deletions(-)
 create mode 100644 idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt
 create mode 100644 idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt
index 87749794e96..bd43e91f610 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt
@@ -17,14 +17,13 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
-import org.jetbrains.kotlin.lexer.JetTokens
+import org.jetbrains.kotlin.idea.core.CommentSaver
 import org.jetbrains.kotlin.psi.JetBlockExpression
 import org.jetbrains.kotlin.psi.JetForExpression
 import org.jetbrains.kotlin.psi.JetPsiFactory
 import org.jetbrains.kotlin.psi.createExpressionByPattern
 import org.jetbrains.kotlin.psi.psiUtil.contentRange
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
-import org.jetbrains.kotlin.psi.psiUtil.getText
 
 public class ConvertToForEachFunctionCallIntention : JetSelfTargetingIntention<JetForExpression>(javaClass(), "Replace with a forEach function call") {
     override fun isApplicableTo(element: JetForExpression, caretOffset: Int): Boolean {
@@ -34,6 +33,8 @@ public class ConvertToForEachFunctionCallIntention : JetSelfTargetingIntention<J
     }
 
     override fun applyTo(element: JetForExpression, editor: Editor) {
+        val commentSaver = CommentSaver(element)
+
         val body = element.getBody()!!
         val loopParameter = element.getLoopParameter()!!
 
@@ -41,6 +42,8 @@ public class ConvertToForEachFunctionCallIntention : JetSelfTargetingIntention<J
 
         val foreachExpression = JetPsiFactory(element).createExpressionByPattern(
                 "$0.forEach{$1->$2}", element.getLoopRange()!!, loopParameter, functionBodyArgument)
-        element.replace(foreachExpression)
+        val result = element.replace(foreachExpression)
+
+        commentSaver.restoreComments(result)
     }
 }
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt b/idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt
new file mode 100644
index 00000000000..105acad5688
--- /dev/null
+++ b/idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt
@@ -0,0 +1,3 @@
+fun foo() {
+    <caret>for (x /* current */ in 1..10/* from 1 to 10 */) { }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt.after
new file mode 100644
index 00000000000..3bd76b4a0c3
--- /dev/null
+++ b/idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt.after
@@ -0,0 +1,3 @@
+fun foo() {
+    (1..10/* from 1 to 10 */).forEach { x /* current */ -> }
+}
\ 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 1c922f6e738..51b4610bac9 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -3943,6 +3943,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("commentsOutsideBody.kt")
+        public void testCommentsOutsideBody() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("emptyBody.kt")
         public void testEmptyBody() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/emptyBody.kt");