Replace with for-each: add new-lines using comment saver

Related to KT-15858

Author: mglukhikh

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

@@ -17,11 +17,8 @@
 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
@@ -40,20 +37,18 @@ class ConvertToForEachFunctionCallIntention : SelfTargetingIntention<KtForExpres
     }
 
     override fun applyTo(element: KtForExpression, editor: Editor?) {
-        val commentSaver = CommentSaver(element)
+        val commentSaver = CommentSaver(element, saveLineBreaks = true)
 
         val labelName = element.getLabelName()
 
         val body = element.body!!
         val loopParameter = element.loopParameter!!
 
-        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 functionBodyArgument: Any = (body as? KtBlockExpression)?.contentRange() ?: body
 
         val psiFactory = KtPsiFactory(element)
         val foreachExpression = psiFactory.createExpressionByPattern(
-            pattern, element.loopRange!!, loopParameter, functionBodyArgument
+            "$0.forEach{$1->\n$2}", element.loopRange!!, loopParameter, functionBodyArgument
         )
         val result = element.replace(foreachExpression) as KtElement
 
@@ -64,14 +59,6 @@ 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/binaryExpressionLoopRange.kt.after

@@ -1,4 +1,6 @@
 // WITH_RUNTIME
 fun main() {
-    1.rangeTo(2).forEach { x -> x }
+    1.rangeTo(2).forEach { x ->
+        x
+    }
 }

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

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

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

@@ -1,4 +1,5 @@
 // WITH_RUNTIME
 fun foo() {
-    (1..10/* from 1 to 10 */).forEach { x /* current */ -> }
+    (1..10/* from 1 to 10 */).forEach { x /* current */ ->
+    }
 }

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

@@ -1,4 +1,5 @@
 // WITH_RUNTIME
 fun foo() {
-    (1..10).forEach { x -> }
+    (1..10).forEach { x ->
+    }
 }

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

@@ -2,5 +2,7 @@
 fun main() {
     val list = 1..4
 
-    list.forEach { x: Int -> x }
+    list.forEach { x: Int ->
+        x
+    }
 }

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

@@ -3,6 +3,8 @@ fun foo() {
     val list = 1..4
     val i = 0
 
-    list.forEach { i -> i }
+    list.forEach { i ->
+        i
+    }
     i
 }

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

@@ -2,5 +2,6 @@
 fun foo() {
     val list = 1..4
 
-    list.forEach { i -> }
+    list.forEach { i ->
+    }
 }

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

@@ -2,5 +2,7 @@
 fun foo() {
     val list = 1..4
 
-    list.forEach { x -> x }
+    list.forEach { x ->
+        x
+    }
 }

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

@@ -2,5 +2,7 @@
 fun main() {
     val list = 1..4
 
-    list.forEach { x: Int -> 11 }
+    list.forEach { x: Int ->
+        11
+    }
 }