Fixed Convert to foreach intention loosing comments

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/jetPsiUtil.kt

@@ -646,11 +646,17 @@ public fun JetElement.getCalleeHighlightingRange(): TextRange {
     return TextRange(startOffset, annotationEntry.getCalleeExpression().endOffset)
 }
 
-public val PsiElement.startOffset: Int
-    get() = getTextRange().getStartOffset()
-
-public val PsiElement.endOffset: Int
-    get() = getTextRange().getEndOffset()
+public fun JetBlockExpression.contentRange(): PsiChildRange {
+    val first = (getLBrace()?.getNextSibling() ?: getFirstChild())
+                        ?.siblings(withItself = false)
+                        ?.firstOrNull { it !is PsiWhiteSpace }
+    val rBrace = getRBrace()
+    if (first == rBrace) return PsiChildRange.EMPTY
+    val last = rBrace!!
+            .siblings(forward = false, withItself = false)
+            .first { it !is PsiWhiteSpace }
+    return PsiChildRange(first, last)
+}
 
 // Annotations on labeled expression lies on it's base expression
 public fun JetExpression.getAnnotationEntries(): List<JetAnnotationEntry> {

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.psi.psiUtil
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+
+//TODO: move here more functions from jetPsiUtil.kt
+
+public val PsiElement.startOffset: Int
+    get() = getTextRange().getStartOffset()
+
+public val PsiElement.endOffset: Int
+    get() = getTextRange().getEndOffset()
+
+public data class PsiChildRange(public val first: PsiElement?, public val last: PsiElement?) : Sequence<PsiElement> {
+    init {
+        if (first == null) {
+            assert(last == null)
+        }
+        else {
+            assert(first.getParent() == last!!.getParent())
+        }
+    }
+
+    public val isEmpty: Boolean
+        get() = first == null
+
+    override fun iterator(): Iterator<PsiElement> {
+        val sequence = if (first == null) {
+            emptySequence<PsiElement>()
+        }
+        else {
+            val afterLast = last!!.getNextSibling()
+            first.siblings().takeWhile { it != afterLast }
+        }
+        return sequence.iterator()
+    }
+
+    companion object {
+        public val EMPTY: PsiChildRange = PsiChildRange(null, null)
+    }
+}
+
+public val PsiElement.allChildren: PsiChildRange
+    get() {
+        val first = getFirstChild()
+        return if (first != null) PsiChildRange(first, getLastChild()) else PsiChildRange.EMPTY
+    }
+
+public val PsiChildRange.textRange: TextRange?
+    get() {
+        if (isEmpty) return null
+        return TextRange(first!!.startOffset, last!!.endOffset)
+    }
+
+public fun PsiChildRange.getText(): String {
+    if (isEmpty) return ""
+    return this.map { it.getText() }.joinToString("")
+}
+

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