From e8904d77273d29b692422a82dead4b1674b42c6e Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Fri, 22 Apr 2016 16:17:19 +0300
Subject: [PATCH] Merging loop with initialization before in more cases +
 better formatting

---
 .../loopToCallChain/baseTransformations.kt    | 24 +++++++++----------
 .../intentions/loopToCallChain/commonUtils.kt |  8 +++++++
 .../idea/intentions/loopToCallChain/utils.kt  |  2 +-
 .../moveInitializationsCloserToLoop.kt        | 15 ++++++++++++
 .../moveInitializationsCloserToLoop.kt.after  | 12 ++++++++++
 5 files changed, 48 insertions(+), 13 deletions(-)
 create mode 100644 idea/testData/intentions/loopToCallChain/moveInitializationsCloserToLoop.kt
 create mode 100644 idea/testData/intentions/loopToCallChain/moveInitializationsCloserToLoop.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/baseTransformations.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/baseTransformations.kt
index e961941fea8..d91a76c8a3c 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/baseTransformations.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/baseTransformations.kt
@@ -42,22 +42,20 @@ abstract class AssignToVariableResultTransformation(
     override fun convertLoop(resultCallChain: KtExpression, commentSavingRangeHolder: CommentSavingRangeHolder): KtExpression {
         initialization.initializer.replace(resultCallChain)
 
-        val loopUnwrapped = loop.unwrapIfLabeled()
-        val previousStatement = loopUnwrapped.previousStatement()
-
-        loopUnwrapped.delete()
-
-        if (initialization.variable.isVar && !initialization.variable.hasWriteUsages()) { // change variable to 'val' if possible
-            initialization.variable.valOrVarKeyword.replace(KtPsiFactory(initialization.variable).createValKeyword())
+        val variable = initialization.variable
+        if (variable.isVar && variable.countWriteUsages() == variable.countWriteUsages(loop)) { // change variable to 'val' if possible
+            variable.valOrVarKeyword.replace(KtPsiFactory(variable).createValKeyword())
         }
 
-        // move initializer to the place where the loop was if needed
+        val loopUnwrapped = loop.unwrapIfLabeled()
+
+        // move initializer to the loop if needed
         var initializationStatement = initialization.initializationStatement
-        if (initializationStatement != previousStatement) {
-            val block = initializationStatement.parent
+        if (initializationStatement.nextStatement() != loopUnwrapped) {
+            val block = loopUnwrapped.parent
             assert(block is KtBlockExpression)
-            val movedInitializationStatement = block.addAfter(initializationStatement, previousStatement) as KtExpression
-            block.addAfter(KtPsiFactory(block).createNewLine(), previousStatement)
+            val movedInitializationStatement = block.addBefore(initializationStatement, loopUnwrapped) as KtExpression
+            block.addBefore(KtPsiFactory(block).createNewLine(), loopUnwrapped)
 
             commentSavingRangeHolder.remove(initializationStatement)
 
@@ -65,6 +63,8 @@ abstract class AssignToVariableResultTransformation(
             initializationStatement = movedInitializationStatement
         }
 
+        loopUnwrapped.delete()
+
         return initializationStatement
     }
 
diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/commonUtils.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/commonUtils.kt
index 389ad459270..d919550d6d6 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/commonUtils.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/commonUtils.kt
@@ -92,6 +92,14 @@ fun KtCallableDeclaration.countUsages(): Int {
     return ReferencesSearch.search(this, useScope).count()
 }
 
+fun KtVariableDeclaration.countWriteUsages(): Int {
+    assert(this.isPhysical)
+    if (!isVar) return 0
+    return ReferencesSearch.search(this, useScope).count {
+        (it as? KtSimpleNameReference)?.element?.readWriteAccess(useResolveForReadWrite = true)?.isWrite == true
+    }
+}
+
 fun KtVariableDeclaration.countWriteUsages(inElement: KtElement): Int {
     assert(this.isPhysical)
     if (!isVar) return 0
diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt
index c8c85b18718..e74af0f675d 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt
@@ -254,7 +254,7 @@ fun canSwapExecutionOrder(expressionBefore: KtExpression, expressionAfter: KtExp
         return canSwapExecutionOrder(expressionBefore, expressionAfter.right ?: return true)
     }
 
-    if (expressionBefore.isConstant() || expressionAfter.isConstant()) return true
+    if (expressionBefore.hasNoSideEffect() || expressionAfter.hasNoSideEffect()) return true
 
     //TODO: more cases
     return false
diff --git a/idea/testData/intentions/loopToCallChain/moveInitializationsCloserToLoop.kt b/idea/testData/intentions/loopToCallChain/moveInitializationsCloserToLoop.kt
new file mode 100644
index 00000000000..edb0b23bd41
--- /dev/null
+++ b/idea/testData/intentions/loopToCallChain/moveInitializationsCloserToLoop.kt
@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'map{}'"
+import java.util.*
+
+fun foo(list: List<String>) {
+    val result = ArrayList<Int>()
+
+    bar()
+
+    <caret>for (s in list) {
+        result.add(s.length)
+    }
+}
+
+fun bar(){}
\ No newline at end of file
diff --git a/idea/testData/intentions/loopToCallChain/moveInitializationsCloserToLoop.kt.after b/idea/testData/intentions/loopToCallChain/moveInitializationsCloserToLoop.kt.after
new file mode 100644
index 00000000000..671ed988343
--- /dev/null
+++ b/idea/testData/intentions/loopToCallChain/moveInitializationsCloserToLoop.kt.after
@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'map{}'"
+import java.util.*
+
+fun foo(list: List<String>) {
+
+    bar()
+
+    val <caret>result = list.map { it.length }
+}
+
+fun bar(){}
\ No newline at end of file