From 46055dfd2dfab091b99282716510f37f02305bc6 Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Fri, 13 May 2016 15:41:00 +0300
Subject: [PATCH] Refactoring

---
 .../intentions/loopToCallChain/interfaces.kt  |  2 +-
 .../loopToCallChain/matchAndConvert.kt        | 55 ++++++++++---------
 2 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt
index 478330a4ca3..17259cd2b09 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt
@@ -103,7 +103,7 @@ data class MatchingState(
          */
         val indexVariable: KtCallableDeclaration?,
         val initializationStatementsToDelete: Collection<KtExpression> = emptyList(),
-        val previousTransformations: List<SequenceTransformation>,
+        val previousTransformations: MutableList<SequenceTransformation> = arrayListOf(),
         val pseudocodeProvider: () -> Pseudocode
 )
 
diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt
index d1dea98415e..ba8f2194b6a 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt
@@ -46,7 +46,6 @@ import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfo
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.utils.addToStdlib.check
-import java.util.*
 
 object MatcherRegistrar {
     val matchers: Collection<TransformationMatcher> = listOf(
@@ -70,31 +69,11 @@ data class MatchResult(
 fun match(loop: KtForExpression): MatchResult? {
     val (inputVariable, indexVariable, sequenceExpression) = extractLoopData(loop) ?: return null
 
+    var state = createInitialMatchingState(loop, inputVariable, indexVariable) ?: return null
+
     // used just as optimization to avoid unnecessary checks
     val loopContainsEmbeddedBreakOrContinue = loop.containsEmbeddedBreakOrContinue()
 
-    val sequenceTransformations = ArrayList<SequenceTransformation>()
-
-    val pseudocodeProvider: () -> Pseudocode = object : () -> Pseudocode {
-        val pseudocode: Pseudocode by lazy {
-            val declaration = loop.containingDeclarationForPseudocode!!
-            val bindingContext = loop.analyze(BodyResolveMode.FULL)
-            PseudocodeUtil.generatePseudocode(declaration, bindingContext)
-        }
-
-        override fun invoke() = pseudocode
-    }
-
-    var state = MatchingState(
-            outerLoop = loop,
-            innerLoop = loop,
-            statements = listOf(loop.body ?: return null),
-            inputVariable = inputVariable,
-            indexVariable = indexVariable,
-            previousTransformations = sequenceTransformations,
-            pseudocodeProvider = pseudocodeProvider
-    )
-
     MatchLoop@
     while (true) {
         state = state.unwrapBlock()
@@ -116,7 +95,7 @@ fun match(loop: KtForExpression): MatchResult? {
             if (match != null) {
                 if (!inputVariableUsed && match.allTransformations.any { it.shouldUseInputVariable }) return null
 
-                sequenceTransformations.addAll(match.sequenceTransformations)
+                state.previousTransformations += match.sequenceTransformations
 
                 when (match) {
                     is TransformationMatch.Sequence -> {
@@ -143,7 +122,7 @@ fun match(loop: KtForExpression): MatchResult? {
                     is TransformationMatch.Result -> {
                         if (restContainsEmbeddedBreakOrContinue && !matcher.embeddedBreakOrContinuePossible) continue@MatchersLoop
 
-                        return TransformationMatch.Result(match.resultTransformation, sequenceTransformations)
+                        return TransformationMatch.Result(match.resultTransformation, state.previousTransformations)
                                 .let { mergeTransformations(it) }
                                 .let { MatchResult(sequenceExpression, it, state.initializationStatementsToDelete) }
                                 .check { checkSmartCastsPreserved(loop, it) }
@@ -209,6 +188,32 @@ private fun extractLoopData(loop: KtForExpression): LoopData? {
     return LoopData(loop.loopParameter ?: return null, null, loopRange)
 }
 
+private fun createInitialMatchingState(
+        loop: KtForExpression,
+        inputVariable: KtCallableDeclaration,
+        indexVariable: KtCallableDeclaration?
+): MatchingState? {
+
+    val pseudocodeProvider: () -> Pseudocode = object : () -> Pseudocode {
+        val pseudocode: Pseudocode by lazy {
+            val declaration = loop.containingDeclarationForPseudocode!!
+            val bindingContext = loop.analyze(BodyResolveMode.FULL)
+            PseudocodeUtil.generatePseudocode(declaration, bindingContext)
+        }
+
+        override fun invoke() = pseudocode
+    }
+
+    return MatchingState(
+            outerLoop = loop,
+            innerLoop = loop,
+            statements = listOf(loop.body ?: return null),
+            inputVariable = inputVariable,
+            indexVariable = indexVariable,
+            pseudocodeProvider = pseudocodeProvider
+    )
+}
+
 private fun isExpressionTypeSupported(expression: KtExpression): Boolean {
     val type = expression.analyze(BodyResolveMode.PARTIAL).getType(expression) ?: return false