Replace with for-each: add new-line if necessary #KT-15858 Fixed

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

@@ -17,8 +17,11 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.refactoring.getLineNumber
 import org.jetbrains.kotlin.idea.util.CommentSaver
+import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.contentRange
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
@@ -26,7 +29,10 @@ import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
 import java.util.*
 
-class ConvertToForEachFunctionCallIntention : SelfTargetingIntention<KtForExpression>(KtForExpression::class.java, "Replace with a 'forEach' function call") {
+class ConvertToForEachFunctionCallIntention : SelfTargetingIntention<KtForExpression>(
+    KtForExpression::class.java,
+    "Replace with a 'forEach' function call"
+) {
     override fun isApplicableTo(element: KtForExpression, caretOffset: Int): Boolean {
         val rParen = element.rightParenthesis ?: return false
         if (caretOffset > rParen.endOffset) return false // available only on the loop header, not in the body
@@ -41,11 +47,14 @@ class ConvertToForEachFunctionCallIntention : SelfTargetingIntention<KtForExpres
         val body = element.body!!
         val loopParameter = element.loopParameter!!
 
-        val functionBodyArgument: Any = (body as? KtBlockExpression)?.contentRange() ?: body
+        val blockExpression = body as? KtBlockExpression
+        val functionBodyArgument: Any = blockExpression?.contentRange() ?: body
+        val pattern = if (needLineBreaks(blockExpression)) "$0.forEach{$1->\n$2\n}" else "$0.forEach{$1->$2}"
 
         val psiFactory = KtPsiFactory(element)
         val foreachExpression = psiFactory.createExpressionByPattern(
-                "$0.forEach{$1->$2}", element.loopRange!!, loopParameter, functionBodyArgument)
+            pattern, element.loopRange!!, loopParameter, functionBodyArgument
+        )
         val result = element.replace(foreachExpression) as KtElement
 
         result.findDescendantOfType<KtFunctionLiteral>()!!.getContinuesWithLabel(labelName).forEach {
@@ -55,6 +64,14 @@ class ConvertToForEachFunctionCallIntention : SelfTargetingIntention<KtForExpres
         commentSaver.restore(result)
     }
 
+    private fun needLineBreaks(blockExpression: KtBlockExpression?): Boolean {
+        val contentRange = blockExpression?.contentRange() ?: return false
+        if ((contentRange.last as? PsiComment)?.tokenType != KtTokens.EOL_COMMENT) return false
+        val statements = blockExpression.statements
+        val fistStatementLine = statements.firstOrNull()?.getLineNumber()
+        return statements.all { it.getLineNumber() == fistStatementLine }
+    }
+
     private fun KtElement.getContinuesWithLabel(labelName: String?): List<KtContinueExpression> {
         val continueElements = ArrayList<KtContinueExpression>()
 

idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    <caret>for (l in listOf(1, 2)) {
+        /* comment */
+    }
+}

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

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    listOf(1, 2).forEach { l -> /* comment */ }
+}

idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment1.kt

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun test() {
+    <caret>for (l in listOf(1, 2)) {
+        // comment
+    }
+}

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

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun test() {
+    listOf(1, 2).forEach { l ->
+        // comment
+    }
+}

idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment2.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    <caret>for (l in listOf(1, 2)) {
+        foo() // comment
+    }
+}

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

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    listOf(1, 2).forEach { l ->
+        foo() // comment
+    }
+}

idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment3.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    <caret>for (l in listOf(1, 2)) {
+        foo(); foo() // comment
+    }
+}

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

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    listOf(1, 2).forEach { l ->
+        foo(); foo() // comment
+    }
+}

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

@@ -6525,6 +6525,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("blockCommentOnly.kt")
+        public void testBlockCommentOnly() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("commentsInBody.kt")
         public void testCommentsInBody() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/commentsInBody.kt");
@@ -6561,6 +6567,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("endOfLineComment1.kt")
+        public void testEndOfLineComment1() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment1.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("endOfLineComment2.kt")
+        public void testEndOfLineComment2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment2.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("endOfLineComment3.kt")
+        public void testEndOfLineComment3() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment3.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("iterativeElementTypeSpecified.kt")
         public void testIterativeElementTypeSpecified() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToForEachFunctionCall/iterativeElementTypeSpecified.kt");