From bec28c8388d050d371c333391bc0b3f81cf35a7a Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <Mikhail.Glukhikh@jetbrains.com>
Date: Wed, 28 Feb 2018 11:25:29 +0300
Subject: [PATCH] Replace with for-each: add new-lines using comment saver

Related to KT-15858
---
 .../ConvertToForEachFunctionCallIntention.kt  | 19 +++----------------
 .../binaryExpressionLoopRange.kt.after        |  4 +++-
 .../blockCommentOnly.kt.after                 |  4 +++-
 .../commentsOutsideBody.kt.after              |  3 ++-
 .../emptyBody.kt.after                        |  3 ++-
 .../iterativeElementTypeSpecified.kt.after    |  4 +++-
 .../noCurlyBraces.kt.after                    |  4 +++-
 .../noStatements.kt.after                     |  3 ++-
 .../simple.kt.after                           |  4 +++-
 .../typeAnnotatedWithNonBlockBody.kt.after    |  4 +++-
 10 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt
index 94532f86672..051526090a0 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt
+++ b/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>()
 
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/binaryExpressionLoopRange.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/binaryExpressionLoopRange.kt.after
index c923c8bb267..e4f9548a9c3 100644
--- a/idea/testData/intentions/convertToForEachFunctionCall/binaryExpressionLoopRange.kt.after
+++ b/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
+    }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt.after
index d7c57129343..927aab6c2d0 100644
--- a/idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt.after
+++ b/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 */
+    }
 }
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt.after
index 5089b7cacb2..76f3d1ea6fc 100644
--- a/idea/testData/intentions/convertToForEachFunctionCall/commentsOutsideBody.kt.after
+++ b/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 */ ->
+    }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/emptyBody.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/emptyBody.kt.after
index 4b05bef6659..9e089bd6399 100644
--- a/idea/testData/intentions/convertToForEachFunctionCall/emptyBody.kt.after
+++ b/idea/testData/intentions/convertToForEachFunctionCall/emptyBody.kt.after
@@ -1,4 +1,5 @@
 // WITH_RUNTIME
 fun foo() {
-    (1..10).forEach { x -> }
+    (1..10).forEach { x ->
+    }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/iterativeElementTypeSpecified.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/iterativeElementTypeSpecified.kt.after
index e800c27fe50..0eb953cbe3a 100644
--- a/idea/testData/intentions/convertToForEachFunctionCall/iterativeElementTypeSpecified.kt.after
+++ b/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
+    }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/noCurlyBraces.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/noCurlyBraces.kt.after
index 641258e5004..37f75ba764c 100644
--- a/idea/testData/intentions/convertToForEachFunctionCall/noCurlyBraces.kt.after
+++ b/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
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/noStatements.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/noStatements.kt.after
index 1462117e2f3..fd378bbb0de 100644
--- a/idea/testData/intentions/convertToForEachFunctionCall/noStatements.kt.after
+++ b/idea/testData/intentions/convertToForEachFunctionCall/noStatements.kt.after
@@ -2,5 +2,6 @@
 fun foo() {
     val list = 1..4
 
-    list.forEach { i -> }
+    list.forEach { i ->
+    }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/simple.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/simple.kt.after
index 2bea915284c..051d560b18f 100644
--- a/idea/testData/intentions/convertToForEachFunctionCall/simple.kt.after
+++ b/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
+    }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/typeAnnotatedWithNonBlockBody.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/typeAnnotatedWithNonBlockBody.kt.after
index 5f8d6656adf..668116ce571 100644
--- a/idea/testData/intentions/convertToForEachFunctionCall/typeAnnotatedWithNonBlockBody.kt.after
+++ b/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
+    }
 }
\ No newline at end of file