Fixed Convert to foreach intention loosing comments

idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt

@@ -17,11 +17,14 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
+import org.jetbrains.kotlin.lexer.JetTokens
 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 {
@@ -35,7 +38,15 @@ public class ConvertToForEachFunctionCallIntention : JetSelfTargetingIntention<J
         val loopParameter = element.getLoopParameter()!!
 
         val functionBodyText = when (body) {
-            is JetBlockExpression -> body.getStatements().map { it.getText() }.joinToString("\n")
+            is JetBlockExpression -> {
+                val content = body.contentRange()
+                val text = content.getText()
+                if (content.last?.getNode()?.getElementType() == JetTokens.EOL_COMMENT)
+                    text + "\n"
+                else
+                    text
+            }
+
             else -> body.getText()
         }
 

idea/src/org/jetbrains/kotlin/idea/intentions/InvertIfConditionIntention.kt

@@ -143,8 +143,8 @@ public class InvertIfConditionIntention : JetSelfTargetingIntention<JetIfExpress
 
         if (thenBranch is JetBlockExpression) {
             val range = thenBranch.contentRange()
-            if (range != null) {
-                parent.addRangeAfter(range.first, range.second, ifExpression)
+            if (!range.isEmpty) {
+                parent.addRangeAfter(range.first, range.last, ifExpression)
                 parent.addAfter(factory.createNewLine(), ifExpression)
             }
         }
@@ -155,14 +155,6 @@ public class InvertIfConditionIntention : JetSelfTargetingIntention<JetIfExpress
         return ifExpression
     }
 
-    private fun JetBlockExpression.contentRange(): Pair<PsiElement, PsiElement>? {
-        val first = getLBrace()?.siblings(withItself = false)?.firstOrNull { it !is PsiWhiteSpace } ?: return null
-        val rBrace = getRBrace()
-        if (first == rBrace) return null
-        val last = rBrace!!.siblings(forward = false, withItself = false).first { it !is PsiWhiteSpace }
-        return Pair(first, last)
-    }
-
     private fun exitStatementExecutedAfter(expression: JetExpression): JetExpression? {
         val parent = expression.getParent()
         if (parent is JetBlockExpression) {

idea/testData/intentions/convertToForEachFunctionCall/commentsInBody.kt

@@ -0,0 +1,11 @@
+fun foo() {
+    val list = 1..4
+
+    <caret>for (x in list) { // start of loop
+        // comment 1
+        var v = x + 1
+        // comment 2
+        v++
+        // end of loop
+    }
+}

idea/testData/intentions/convertToForEachFunctionCall/commentsInBody.kt.after

@@ -0,0 +1,11 @@
+fun foo() {
+    val list = 1..4
+
+    list.forEach { x ->// start of loop
+        // comment 1
+        var v = x + 1
+        // comment 2
+        v++
+        // end of loop
+    }
+}

idea/testData/intentions/convertToForEachFunctionCall/commentsInBody2.kt

@@ -0,0 +1,8 @@
+fun foo() {
+    val list = 1..4
+
+    <caret>for (x in list) {
+        // comment
+        var v = x + 1
+    }
+}

idea/testData/intentions/convertToForEachFunctionCall/commentsInBody2.kt.after

@@ -0,0 +1,7 @@
+fun foo() {
+    val list = 1..4
+
+    list.forEach { x ->// comment
+        var v = x + 1
+    }
+}

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

@@ -3898,6 +3898,18 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("commentsInBody.kt")
+        public void testCommentsInBody() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/commentsInBody.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("commentsInBody2.kt")
+        public void testCommentsInBody2() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/commentsInBody2.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("iterativeElementTypeSpecified.kt")
         public void testIterativeElementTypeSpecified() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/iterativeElementTypeSpecified.kt");