Merging loop with initialization before in more cases + better formatting

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 variable = initialization.variable
-        val previousStatement = loopUnwrapped.previousStatement()
+        if (variable.isVar && variable.countWriteUsages() == variable.countWriteUsages(loop)) { // change variable to 'val' if possible
-
+            variable.valOrVarKeyword.replace(KtPsiFactory(variable).createValKeyword())
-        loopUnwrapped.delete()
-
-        if (initialization.variable.isVar && !initialization.variable.hasWriteUsages()) { // change variable to 'val' if possible
-            initialization.variable.valOrVarKeyword.replace(KtPsiFactory(initialization.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) {
+        if (initializationStatement.nextStatement() != loopUnwrapped) {
-            val block = initializationStatement.parent
+            val block = loopUnwrapped.parent
             assert(block is KtBlockExpression)
-            val movedInitializationStatement = block.addAfter(initializationStatement, previousStatement) as KtExpression
+            val movedInitializationStatement = block.addBefore(initializationStatement, loopUnwrapped) as KtExpression
-            block.addAfter(KtPsiFactory(block).createNewLine(), previousStatement)
+            block.addBefore(KtPsiFactory(block).createNewLine(), loopUnwrapped)
 
             commentSavingRangeHolder.remove(initializationStatement)
 
@@ -65,6 +63,8 @@ abstract class AssignToVariableResultTransformation(
             initializationStatement = movedInitializationStatement
         }
 
+        loopUnwrapped.delete()
+
         return initializationStatement
     }
 

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

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

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

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