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>()