Merged all map-transformations into one class

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

@@ -49,19 +49,7 @@ class AddToCollectionTransformation(
             }
 
             is MapTransformation -> {
-                MapToTransformation.create(loop, previousTransformation.inputVariable, null, targetCollection, previousTransformation.mapping, mapNotNull = false)
+                MapToTransformation.create(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, targetCollection, previousTransformation.mapping, previousTransformation.mapNotNull)
-            }
-
-            is MapNotNullTransformation -> {
-                MapToTransformation.create(loop, previousTransformation.inputVariable, null, targetCollection, previousTransformation.mapping, mapNotNull = true)
-            }
-
-            is MapIndexedTransformation -> {
-                MapToTransformation.create(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, targetCollection, previousTransformation.mapping, mapNotNull = false)
-            }
-
-            is MapIndexedNotNullTransformation -> {
-                MapToTransformation.create(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, targetCollection, previousTransformation.mapping, mapNotNull = true)
             }
 
             is FlatMapTransformation -> {
@@ -144,7 +132,7 @@ class AddToCollectionTransformation(
                                 AssignToListTransformation(state.outerLoop, collectionInitialization)
                             }
                             else {
-                                val mapTransformation = MapTransformation(state.outerLoop, state.inputVariable, addOperationArgument)
+                                val mapTransformation = MapTransformation(state.outerLoop, state.inputVariable, null, addOperationArgument, mapNotNull = false)
                                 AssignSequenceTransformationResultTransformation(mapTransformation, collectionInitialization)
                             }
                             return ResultTransformationMatch(transformation)
@@ -176,7 +164,7 @@ class AddToCollectionTransformation(
                         return ResultTransformationMatch(assignToSetTransformation)
                     }
                     else {
-                        val mapTransformation = MapTransformation(state.outerLoop, state.inputVariable, addOperationArgument)
+                        val mapTransformation = MapTransformation(state.outerLoop, state.inputVariable, null, addOperationArgument, mapNotNull = false)
                         return ResultTransformationMatch(assignToSetTransformation, mapTransformation)
                     }
                 }

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

@@ -164,11 +164,10 @@ class FilterIsInstanceTransformation(
 
 class FilterNotNullTransformation(override val loop: KtForExpression) : SequenceTransformation {
     override fun mergeWithPrevious(previousTransformation: SequenceTransformation): SequenceTransformation? {
-        return when (previousTransformation) {
+        if (previousTransformation is MapTransformation) {
-            is MapTransformation -> MapNotNullTransformation(loop, previousTransformation.inputVariable, previousTransformation.mapping)
+            return MapTransformation(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, previousTransformation.mapping, mapNotNull = true)
-            is MapIndexedTransformation -> MapIndexedNotNullTransformation(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, previousTransformation.mapping)
-            else -> null
         }
+        return null
     }
 
     override val affectsIndex: Boolean

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

@@ -65,7 +65,7 @@ class FlatMapTransformation(
 
             if (state.indexVariable != null && state.indexVariable.hasUsages(transform)) {
                 // if nested loop range uses index, convert to "mapIndexed {...}.flatMap { it }"
-                val mapIndexedTransformation = MapIndexedTransformation(state.outerLoop, state.inputVariable, state.indexVariable, transform)
+                val mapIndexedTransformation = MapTransformation(state.outerLoop, state.inputVariable, state.indexVariable, transform, mapNotNull = false)
                 val inputVarExpression = KtPsiFactory(nestedLoop).createExpressionByPattern("$0", state.inputVariable.nameAsSafeName)
                 val flatMapTransformation = FlatMapTransformation(state.outerLoop, state.inputVariable, inputVarExpression)
                 val newState = state.copy(

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

@@ -23,18 +23,28 @@ import org.jetbrains.kotlin.psi.*
 class MapTransformation(
         override val loop: KtForExpression,
         val inputVariable: KtCallableDeclaration,
-        val mapping: KtExpression
+        val indexVariable: KtCallableDeclaration?,
+        val mapping: KtExpression,
+        val mapNotNull: Boolean
 ) : SequenceTransformation {
 
+    private val functionName = if (indexVariable != null)
+        if (mapNotNull) "mapIndexedNotNull" else "mapIndexed"
+    else
+        if (mapNotNull) "mapNotNull" else "map"
+
     override val affectsIndex: Boolean
-        get() = false
+        get() = mapNotNull
 
     override val presentation: String
-        get() = "map{}"
+        get() = "$functionName{}"
 
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
-        val lambda = generateLambda(inputVariable, mapping)
+        val lambda = if (indexVariable != null)
-        return chainedCallGenerator.generate("map$0:'{}'", lambda)
+            generateLambda(mapping, indexVariable, inputVariable)
+        else
+            generateLambda(inputVariable, mapping)
+        return chainedCallGenerator.generate("$functionName$0:'{}'", lambda)
     }
 
     /**
@@ -62,9 +72,9 @@ class MapTransformation(
                 if (mapping.containsEmbeddedBreakOrContinue()) return null
 
                 val transformation = if (state.indexVariable != null && state.indexVariable.hasUsages(mapping))
-                    MapIndexedNotNullTransformation(state.outerLoop, state.inputVariable, state.indexVariable, mapping)
+                    MapTransformation(state.outerLoop, state.inputVariable, state.indexVariable, mapping, mapNotNull = true)
                 else
-                    MapNotNullTransformation(state.outerLoop, state.inputVariable, mapping)
+                    MapTransformation(state.outerLoop, state.inputVariable, null, mapping, mapNotNull = true)
                 val newState = state.copy(statements = restStatements, inputVariable = declaration)
                 return SequenceTransformationMatch(transformation, newState)
             }
@@ -72,67 +82,11 @@ class MapTransformation(
             if (initializer.containsEmbeddedBreakOrContinue()) return null
 
             val transformation = if (state.indexVariable != null && state.indexVariable.hasUsages(initializer))
-                MapIndexedTransformation(state.outerLoop, state.inputVariable, state.indexVariable, initializer)
+                MapTransformation(state.outerLoop, state.inputVariable, state.indexVariable, initializer, mapNotNull = false)
             else
-                MapTransformation(state.outerLoop, state.inputVariable, initializer)
+                MapTransformation(state.outerLoop, state.inputVariable, null, initializer, mapNotNull = false)
             val newState = state.copy(statements = restStatements, inputVariable = declaration)
             return SequenceTransformationMatch(transformation, newState)
         }
     }
 }
-
-class MapIndexedTransformation(
-        override val loop: KtForExpression,
-        val inputVariable: KtCallableDeclaration,
-        val indexVariable: KtCallableDeclaration,
-        val mapping: KtExpression
-) : SequenceTransformation {
-
-    override val affectsIndex: Boolean
-        get() = false
-
-    override val presentation: String
-        get() = "mapIndexed{}"
-
-    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
-        val lambda = generateLambda(mapping, indexVariable, inputVariable)
-        return chainedCallGenerator.generate("mapIndexed $0:'{}'", lambda)
-    }
-}
-
-class MapNotNullTransformation(
-        override val loop: KtForExpression,
-        val inputVariable: KtCallableDeclaration,
-        val mapping: KtExpression
-) : SequenceTransformation {
-
-    override val affectsIndex: Boolean
-        get() = true
-
-    override val presentation: String
-        get() = "mapNotNull{}"
-
-    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
-        val lambda = generateLambda(inputVariable, mapping)
-        return chainedCallGenerator.generate("mapNotNull$0:'{}'", lambda)
-    }
-}
-
-class MapIndexedNotNullTransformation(
-        override val loop: KtForExpression,
-        val inputVariable: KtCallableDeclaration,
-        val indexVariable: KtCallableDeclaration,
-        val mapping: KtExpression
-) : SequenceTransformation {
-
-    override val affectsIndex: Boolean
-        get() = false
-
-    override val presentation: String
-        get() = "mapIndexedNotNull{}"
-
-    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
-        val lambda = generateLambda(mapping, indexVariable, inputVariable)
-        return chainedCallGenerator.generate("mapIndexedNotNull $0:'{}'", lambda)
-    }
-}