Supported case when result variable initialization is not right before the loop

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/baseTransformations.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain
 
 import org.jetbrains.kotlin.idea.core.replaced
+import org.jetbrains.kotlin.psi.KtBlockExpression
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtForExpression
 import org.jetbrains.kotlin.psi.KtPsiFactory
@@ -40,19 +41,34 @@ abstract class AssignToVariableResultTransformation(
 
     override val commentSavingRange = PsiChildRange(initialization.initializationStatement, loop.unwrapIfLabeled())
 
+    //TODO: does not work in case of move
     private val commentRestoringRange = commentSavingRange.withoutLastStatement()
 
     override fun commentRestoringRange(convertLoopResult: KtExpression) = commentRestoringRange
 
     override fun convertLoop(resultCallChain: KtExpression): KtExpression {
         initialization.initializer.replace(resultCallChain)
+
+        val previousStatement = loop.unwrapIfLabeled().previousStatement()
+
         loop.deleteWithLabels()
 
         if (initialization.variable.isVar && !initialization.variable.hasWriteUsages()) { // change variable to 'val' if possible
             initialization.variable.valOrVarKeyword.replace(KtPsiFactory(initialization.variable).createValKeyword())
         }
 
-        return initialization.initializationStatement
+        // move initializer to the place where the loop was if needed
+        var initializationStatement = initialization.initializationStatement
+        if (initializationStatement != previousStatement) {
+            val block = initializationStatement.parent
+            assert(block is KtBlockExpression)
+            val movedInitializationStatement = block.addAfter(initializationStatement, previousStatement) as KtExpression
+            block.addAfter(KtPsiFactory(block).createNewLine(), previousStatement)
+            initializationStatement.delete()
+            initializationStatement = movedInitializationStatement
+        }
+
+        return initializationStatement
     }
 }
 

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/commonUtils.kt

@@ -32,10 +32,16 @@ import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelector
 import org.jetbrains.kotlin.psi.psiUtil.parents
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
+fun KtExpression.isConstant(): Boolean {
+    val bindingContext = analyze(BodyResolveMode.PARTIAL)
+    return ConstantExpressionEvaluator.getConstant(this, bindingContext) != null
+}
+
 fun KtExpression?.isTrueConstant()
         = this != null && node?.elementType == KtNodeTypes.BOOLEAN_CONSTANT && text == "true"
 

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt

@@ -17,6 +17,8 @@
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain
 
 import com.intellij.openapi.util.Key
+import com.intellij.psi.PsiDocumentManager
+import com.intellij.psi.codeStyle.CodeStyleManager
 import org.jetbrains.kotlin.idea.analysis.analyzeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
@@ -149,6 +151,14 @@ fun convertLoop(loop: KtForExpression, matchResult: MatchResult): KtExpression {
 
     commentSaver.restore(resultTransformation.commentRestoringRange(result))
 
+    // need to manually adjust indent of the result because in some cases it's made incorrect when moving closer to the loop
+    // TODO: use forceAdjustIndent = true of CommentSaver.restore
+    val file = result.containingFile
+    val psiDocumentManager = PsiDocumentManager.getInstance(file.project)
+    psiDocumentManager.doPostponedOperationsAndUnblockDocument(psiDocumentManager.getDocument(file)!!)
+    val codeStyleManager = CodeStyleManager.getInstance(file.project)
+    codeStyleManager.adjustLineIndent(file, result.textRange)
+
     return result
 }
 

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndAssignTransformation.kt

@@ -16,7 +16,6 @@
 
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain.result
 
-import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.*
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterTransformation
 import org.jetbrains.kotlin.lexer.KtTokens
@@ -24,8 +23,6 @@ import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtBreakExpression
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtForExpression
-import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
-import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
 class FindAndAssignTransformation(
         loop: KtForExpression,
@@ -95,9 +92,7 @@ class FindAndAssignTransformation(
             if (initialization.variable.countUsages(state.outerLoop) != 1) return null // this should be the only usage of this variable inside the loop
 
             // we do not try to convert anything if the initializer is not compile-time constant because of possible side-effects
-            val initializerIsConstant = ConstantExpressionEvaluator.getConstant(
+            if (!initialization.initializer.isConstant()) return null
-                    initialization.initializer, initialization.initializer.analyze(BodyResolveMode.PARTIAL)) != null
-            if (!initializerIsConstant) return null
 
             val generator = buildFindOperationGenerator(right, initialization.initializer, state.inputVariable, findFirst) ?: return null
 

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.idea.imports.importableFqName
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.KtPsiUtil.isOrdinaryAssignment
 import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DelegatingBindingTrace
@@ -36,6 +37,8 @@ import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.utils.addToStdlib.check
+import java.util.*
 
 fun generateLambda(inputVariable: KtCallableDeclaration, expression: KtExpression): KtLambdaExpression {
     val psiFactory = KtPsiFactory(expression)
@@ -99,17 +102,38 @@ fun KtExpression.detectInitializationBeforeLoop(
     if (this !is KtNameReferenceExpression) return null
     if (getQualifiedExpressionForSelector() != null) return null
     val variable = this.mainReference.resolve() as? KtProperty ?: return null
-    val statementBeforeLoop = loop.previousStatement() //TODO: support initialization not right before the loop
 
     // do not allow any other usages of this variable inside the loop
     if (checkNoOtherUsagesInLoop && variable.countUsages(loop) > 1) return null
 
-    if (statementBeforeLoop == variable) {
+    val unwrapped = loop.unwrapIfLabeled()
+    if (unwrapped.parent !is KtBlockExpression) return null
+    val prevStatements = unwrapped
+            .siblings(forward = false, withItself = false)
+            .filterIsInstance<KtExpression>()
+
+    val statementsBetween = ArrayList<KtExpression>()
+    for (statement in prevStatements) {
+        val variableInitialization = extractVariableInitialization(statement, variable)
+        if (variableInitialization != null) {
+            return variableInitialization.check {
+                statementsBetween.all { canSwapExecutionOrder(variableInitialization.initializationStatement, it) }
+            }
+        }
+
+        statementsBetween.add(statement)
+    }
+
+    return null
+}
+
+private fun extractVariableInitialization(statement: KtExpression, variable: KtProperty): VariableInitialization? {
+    if (statement == variable) {
         val initializer = variable.initializer ?: return null
         return VariableInitialization(variable, variable, initializer)
     }
 
-    val assignment = statementBeforeLoop?.asAssignment() ?: return null
+    val assignment = statement.asAssignment() ?: return null
     if (!assignment.left.isVariableReference(variable)) return null
 
     val initializer = assignment.right ?: return null
@@ -197,3 +221,39 @@ fun KtExpression.hasNoSideEffect(): Boolean {
     val classFqName = classDescriptor.importableFqName?.asString()
     return classFqName in NO_SIDE_EFFECT_STANDARD_CLASSES
 }
+
+//TODO: we need more correctness checks (if variable is non-local or is local but can be changed by some local functions)
+fun canSwapExecutionOrder(expressionBefore: KtExpression, expressionAfter: KtExpression): Boolean {
+    assert(expressionBefore.isPhysical)
+    assert(expressionAfter.isPhysical)
+
+    if (expressionBefore is KtDeclaration) {
+        if (expressionBefore !is KtProperty) return false // local function, class or destructuring declaration - do not bother to handle these rare cases
+        if (expressionBefore.hasUsages(expressionAfter)) return false
+        return canSwapExecutionOrder(expressionBefore.initializer ?: return true, expressionAfter)
+    }
+
+    if (expressionAfter is KtDeclaration) {
+        if (expressionAfter !is KtProperty) return false // local function, class or destructuring declaration - do not bother to handle these rare cases
+        return canSwapExecutionOrder(expressionBefore, expressionAfter.initializer ?: return true)
+    }
+
+    if (expressionBefore is KtBinaryExpression && isOrdinaryAssignment(expressionBefore)) {
+        val leftName = expressionBefore.left as? KtSimpleNameExpression ?: return false
+        val target = leftName.mainReference.resolve() as? KtProperty ?: return false
+        if (target.hasUsages(expressionAfter)) return false
+        return canSwapExecutionOrder(expressionBefore.right ?: return true, expressionAfter)
+    }
+
+    if (expressionAfter is KtBinaryExpression && isOrdinaryAssignment(expressionAfter)) {
+        val leftName = expressionAfter.left as? KtSimpleNameExpression ?: return false
+        val target = leftName.mainReference.resolve() as? KtProperty ?: return false
+        if (target.hasUsages(expressionBefore)) return false
+        return canSwapExecutionOrder(expressionBefore, expressionAfter.right ?: return true)
+    }
+
+    if (expressionBefore.isConstant() || expressionAfter.isConstant()) return true
+
+    //TODO: more cases
+    return false
+}

idea/testData/intentions/loopToCallChain/any_ifAssign_moveDeclaration.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'any{}'"
+fun foo(list: List<String>) {
+    var found = false
+    println("Starting the search")
+    <caret>for (s in list) {
+        if (s.length > 0) {
+            found = true
+            break
+        }
+    }
+}

idea/testData/intentions/loopToCallChain/any_ifAssign_moveDeclaration.kt.after

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'any{}'"
+fun foo(list: List<String>) {
+    println("Starting the search")
+    val <caret>found = list.any { it.length > 0 }
+}

idea/testData/intentions/loopToCallChain/any_ifAssign_moveInitialization.kt

@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'any{}'"
+fun foo(list: List<String>, p: Int) {
+    var found: Boolean
+    if (p > 0) {
+        found = false
+        println("Starting the search")
+        <caret>for (s in list) {
+            if (s.length > 0) {
+                found = true
+                break
+            }
+        }
+    }
+    else {
+        found = true
+    }
+}

idea/testData/intentions/loopToCallChain/any_ifAssign_moveInitialization.kt.after

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'any{}'"
+fun foo(list: List<String>, p: Int) {
+    var found: Boolean
+    if (p > 0) {
+        println("Starting the search")
+        <caret>found = list.any { it.length > 0 }
+    }
+    else {
+        found = true
+    }
+}

idea/testData/intentions/loopToCallChain/twoInitializationsBeforeLoop.kt

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterIndexed{}.map{}'"
+import java.util.*
+
+fun foo(list: List<String>): List<Int> {
+    val result = ArrayList<Int>()
+    var i = 0
+    <caret>for (s in list) {
+        if (s.length > i) {
+            result.add(s.length)
+        }
+        i++
+    }
+    return result
+}

idea/testData/intentions/loopToCallChain/twoInitializationsBeforeLoop.kt.after

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterIndexed{}.map{}'"
+import java.util.*
+
+fun foo(list: List<String>): List<Int> {
+    val <caret>result = list
+            .filterIndexed { i, s -> s.length > i }
+            .map { it.length }
+    return result
+}