From 7f880bf58c747974592bf5124cc2e301f0451797 Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Fri, 16 Feb 2018 08:39:07 +0300
Subject: [PATCH] Replace with for-each: add new-line if necessary #KT-15858
 Fixed

---
 .../ConvertToForEachFunctionCallIntention.kt  | 23 +++++++++++++++---
 .../blockCommentOnly.kt                       |  8 +++++++
 .../blockCommentOnly.kt.after                 |  6 +++++
 .../endOfLineComment1.kt                      |  6 +++++
 .../endOfLineComment1.kt.after                |  6 +++++
 .../endOfLineComment2.kt                      |  8 +++++++
 .../endOfLineComment2.kt.after                |  8 +++++++
 .../endOfLineComment3.kt                      |  8 +++++++
 .../endOfLineComment3.kt.after                |  8 +++++++
 .../intentions/IntentionTestGenerated.java    | 24 +++++++++++++++++++
 10 files changed, 102 insertions(+), 3 deletions(-)
 create mode 100644 idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt
 create mode 100644 idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt.after
 create mode 100644 idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment1.kt
 create mode 100644 idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment1.kt.after
 create mode 100644 idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment2.kt
 create mode 100644 idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment2.kt.after
 create mode 100644 idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment3.kt
 create mode 100644 idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment3.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt
index ef7ad6bbef6..94532f86672 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToForEachFunctionCallIntention.kt
+++ b/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>()
 
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt b/idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt
new file mode 100644
index 00000000000..ee40ed3ae47
--- /dev/null
+++ b/idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt
@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    <caret>for (l in listOf(1, 2)) {
+        /* comment */
+    }
+}
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt.after
new file mode 100644
index 00000000000..d7c57129343
--- /dev/null
+++ b/idea/testData/intentions/convertToForEachFunctionCall/blockCommentOnly.kt.after
@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    listOf(1, 2).forEach { l -> /* comment */ }
+}
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment1.kt b/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment1.kt
new file mode 100644
index 00000000000..96ebe3e3625
--- /dev/null
+++ b/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment1.kt
@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun test() {
+    <caret>for (l in listOf(1, 2)) {
+        // comment
+    }
+}
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment1.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment1.kt.after
new file mode 100644
index 00000000000..e19b46feea9
--- /dev/null
+++ b/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment1.kt.after
@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun test() {
+    listOf(1, 2).forEach { l ->
+        // comment
+    }
+}
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment2.kt b/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment2.kt
new file mode 100644
index 00000000000..27c6d245f4a
--- /dev/null
+++ b/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
+    }
+}
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment2.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment2.kt.after
new file mode 100644
index 00000000000..bbe1ed43ff2
--- /dev/null
+++ b/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment2.kt.after
@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo() {}
+
+fun test() {
+    listOf(1, 2).forEach { l ->
+        foo() // comment
+    }
+}
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment3.kt b/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment3.kt
new file mode 100644
index 00000000000..9c9eb78bd22
--- /dev/null
+++ b/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
+    }
+}
diff --git a/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment3.kt.after b/idea/testData/intentions/convertToForEachFunctionCall/endOfLineComment3.kt.after
new file mode 100644
index 00000000000..c2fdf5e9bc1
--- /dev/null
+++ b/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
+    }
+}
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index 1bded1d5ac6..e405d8c704e 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/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");