Minor changes on code review

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

@@ -52,6 +52,8 @@ interface Transformation {
             presentation
     }
 
+    fun mergeWithPrevious(previousTransformation: SequenceTransformation): Transformation?
+
     fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression
 
     val chainCallCount: Int
@@ -65,7 +67,7 @@ interface Transformation {
  * Represents a transformation of input sequence into another sequence
  */
 interface SequenceTransformation : Transformation {
-    fun mergeWithPrevious(previousTransformation: SequenceTransformation): SequenceTransformation? = null
+    override fun mergeWithPrevious(previousTransformation: SequenceTransformation): SequenceTransformation? = null
 
     val affectsIndex: Boolean
 }
@@ -74,7 +76,7 @@ interface SequenceTransformation : Transformation {
  * Represents a final transformation of sequence which produces the result of the whole loop (for example, assigning a found value into a variable).
  */
 interface ResultTransformation : Transformation {
-    fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? = null
+    override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? = null
 
     val commentSavingRange: PsiChildRange
 

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

@@ -287,7 +287,7 @@ private fun MatchResult.generateCallChain(loop: KtForExpression): KtExpression {
 }
 
 private fun mergeTransformations(match: TransformationMatch.Result): TransformationMatch.Result {
-    val transformations = ArrayList<Transformation>().apply { addAll(match.sequenceTransformations); add(match.resultTransformation) }
+    val transformations = (match.sequenceTransformations + match.resultTransformation).toMutableList()
 
     var anyChange: Boolean
     do {
@@ -295,12 +295,7 @@ private fun mergeTransformations(match: TransformationMatch.Result): Transformat
         for (index in 0..transformations.lastIndex - 1) {
             val transformation = transformations[index] as SequenceTransformation
             val next = transformations[index + 1]
-            val merged = when (next) {
+            val merged = next.mergeWithPrevious(transformation) ?: continue
-                is SequenceTransformation -> next.mergeWithPrevious(transformation)
-                is ResultTransformation -> next.mergeWithPrevious(transformation)
-                else -> error("Unknown transformation type: $next")
-            } ?: continue
-
             transformations[index] = merged
             transformations.removeAt(index + 1)
             anyChange = true

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

@@ -122,7 +122,7 @@ class AddToCollectionTransformation(
                 targetCollection: KtExpression,
                 addOperationArgument: KtExpression
         ): TransformationMatch.Result? {
-            val collectionInitialization = targetCollection.detectInitializationBeforeLoop(state.outerLoop, checkNoOtherUsagesInLoop = true) ?: return null
+            val collectionInitialization = targetCollection.isVariableInitializedBeforeLoop(state.outerLoop, checkNoOtherUsagesInLoop = true) ?: return null
             val collectionKind = collectionInitialization.initializer.isSimpleCollectionInstantiation() ?: return null
             val argumentIsInputVariable = addOperationArgument.isVariableReference(state.inputVariable)
             when (collectionKind) {
@@ -223,7 +223,7 @@ class FilterToTransformation private constructor(
                 condition: KtExpression,
                 isInverse: Boolean
         ): ResultTransformation {
-            val initialization = targetCollection.detectInitializationBeforeLoop(loop, checkNoOtherUsagesInLoop = true)
+            val initialization = targetCollection.isVariableInitializedBeforeLoop(loop, checkNoOtherUsagesInLoop = true)
             if (initialization != null && initialization.initializer.hasNoSideEffect()) {
                 val transformation = FilterToTransformation(loop, inputVariable, indexVariable, initialization.initializer, condition, isInverse)
                 return AssignToVariableResultTransformation.createDelegated(transformation, initialization)
@@ -252,7 +252,7 @@ class FilterNotNullToTransformation private constructor(
                 loop: KtForExpression,
                 targetCollection: KtExpression
         ): ResultTransformation {
-            val initialization = targetCollection.detectInitializationBeforeLoop(loop, checkNoOtherUsagesInLoop = true)
+            val initialization = targetCollection.isVariableInitializedBeforeLoop(loop, checkNoOtherUsagesInLoop = true)
             if (initialization != null && initialization.initializer.hasNoSideEffect()) {
                 val transformation = FilterNotNullToTransformation(loop, initialization.initializer)
                 return AssignToVariableResultTransformation.createDelegated(transformation, initialization)
@@ -298,7 +298,7 @@ class MapToTransformation private constructor(
                 mapping: KtExpression,
                 mapNotNull: Boolean
         ): ResultTransformation {
-            val initialization = targetCollection.detectInitializationBeforeLoop(loop, checkNoOtherUsagesInLoop = true)
+            val initialization = targetCollection.isVariableInitializedBeforeLoop(loop, checkNoOtherUsagesInLoop = true)
             if (initialization != null && initialization.initializer.hasNoSideEffect()) {
                 val transformation = MapToTransformation(loop, inputVariable, indexVariable, initialization.initializer, mapping, mapNotNull)
                 return AssignToVariableResultTransformation.createDelegated(transformation, initialization)
@@ -332,7 +332,7 @@ class FlatMapToTransformation private constructor(
                 targetCollection: KtExpression,
                 transform: KtExpression
         ): ResultTransformation {
-            val initialization = targetCollection.detectInitializationBeforeLoop(loop, checkNoOtherUsagesInLoop = true)
+            val initialization = targetCollection.isVariableInitializedBeforeLoop(loop, checkNoOtherUsagesInLoop = true)
             if (initialization != null && initialization.initializer.hasNoSideEffect()) {
                 val transformation = FlatMapToTransformation(loop, inputVariable, initialization.initializer, transform)
                 return AssignToVariableResultTransformation.createDelegated(transformation, initialization)

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

@@ -74,7 +74,7 @@ class CountTransformation(
 
         override fun match(state: MatchingState): TransformationMatch.Result? {
             val operand = state.statements.singleOrNull()?.isPlusPlusOf() ?: return null
-            val initialization = operand.detectInitializationBeforeLoop(state.outerLoop, checkNoOtherUsagesInLoop = true) ?: return null
+            val initialization = operand.isVariableInitializedBeforeLoop(state.outerLoop, checkNoOtherUsagesInLoop = true) ?: return null
 
             if (initialization.variable.countUsages(state.outerLoop) != 1) return null // this should be the only usage of this variable inside the loop
 

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

@@ -78,7 +78,7 @@ object FindTransformationMatcher : TransformationMatcher {
         val left = binaryExpression.left ?: return null
         val right = binaryExpression.right ?: return null
 
-        val initialization = left.detectInitializationBeforeLoop(state.outerLoop, checkNoOtherUsagesInLoop = true) ?: return null
+        val initialization = left.isVariableInitializedBeforeLoop(state.outerLoop, checkNoOtherUsagesInLoop = true) ?: return null
 
         if (initialization.variable.countUsages(state.outerLoop) != 1) return null // this should be the only usage of this variable inside the loop
 

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/IntroduceIndexMatcher.kt

@@ -39,7 +39,7 @@ object IntroduceIndexMatcher : TransformationMatcher {
         // there should be no continuation of the loop in statements before index increment
         if (restStatements.any { statement -> statement.anyDescendantOfType<KtContinueExpression>(::isContinueOfThisLoopOrOuter) }) return null
 
-        val variableInitialization = operand.detectInitializationBeforeLoop(state.outerLoop, checkNoOtherUsagesInLoop = false)
+        val variableInitialization = operand.isVariableInitializedBeforeLoop(state.outerLoop, checkNoOtherUsagesInLoop = false)
                                      ?: return null
         if ((variableInitialization.initializer as? KtConstantExpression)?.text != "0") return null
 

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

@@ -97,7 +97,7 @@ data class VariableInitialization(
         val initializer: KtExpression)
 
 //TODO: we need more correctness checks (if variable is non-local or is local but can be changed by some local functions)
-fun KtExpression.detectInitializationBeforeLoop(
+fun KtExpression.isVariableInitializedBeforeLoop(
         loop: KtForExpression,
         checkNoOtherUsagesInLoop: Boolean
 ): VariableInitialization? {
@@ -147,6 +147,7 @@ enum class CollectionKind {
 }
 
 fun KtExpression.isSimpleCollectionInstantiation(): CollectionKind? {
+    //TODO: support mutableListOf() etc
     val callExpression = this as? KtCallExpression ?: return null //TODO: it can be qualified too
     if (callExpression.valueArguments.isNotEmpty()) return null
     val bindingContext = callExpression.analyze(BodyResolveMode.PARTIAL)