Supported mapIndexedTo and mapIndexedNotNullTo

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

@@ -49,11 +49,19 @@ class AddToCollectionTransformation(
             }
 
             is MapTransformation -> {
-                MapToTransformation.create(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.mapping, mapNotNull = false)
+                MapToTransformation.create(loop, previousTransformation.inputVariable, null, targetCollection, previousTransformation.mapping, mapNotNull = false)
             }
 
             is MapNotNullTransformation -> {
-                MapToTransformation.create(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.mapping, mapNotNull = true)
+                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 -> {
@@ -114,13 +122,9 @@ class AddToCollectionTransformation(
             if (state.indexVariable == null && argumentValue.isVariableReference(state.inputVariable)) {
                 return ResultTransformationMatch(AddToCollectionTransformation(state.outerLoop, targetCollection))
             }
-            else if (state.indexVariable != null) {
-                val mapIndexedTransformation = MapIndexedTransformation(state.outerLoop, state.inputVariable, state.indexVariable, argumentValue)
-                val addToCollectionTransformation = AddToCollectionTransformation(state.outerLoop, targetCollection)
-                return ResultTransformationMatch(addToCollectionTransformation, mapIndexedTransformation)
-            }
             else {
-                return ResultTransformationMatch(MapToTransformation.create(state.outerLoop, state.inputVariable, targetCollection, argumentValue, mapNotNull = false))
+                return ResultTransformationMatch(MapToTransformation.create(
+                        state.outerLoop, state.inputVariable, state.indexVariable, targetCollection, argumentValue, mapNotNull = false))
             }
         }
 
@@ -299,18 +303,25 @@ class FilterNotNullToTransformation private constructor(
 class MapToTransformation private constructor(
         loop: KtForExpression,
         private val inputVariable: KtCallableDeclaration,
+        private val indexVariable: KtCallableDeclaration?,
         private val targetCollection: KtExpression,
         private val mapping: KtExpression,
         mapNotNull: Boolean
 ) : ReplaceLoopResultTransformation(loop) {
 
-    private val functionName = if (mapNotNull) "mapNotNullTo" else "mapTo"
+    private val functionName = if (indexVariable != null)
+        if (mapNotNull) "mapIndexedNotNullTo" else "mapIndexedTo"
+    else
+        if (mapNotNull) "mapNotNullTo" else "mapTo"
 
     override val presentation: String
         get() = "$functionName(){}"
 
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
-        val lambda = generateLambda(inputVariable, mapping)
+        val lambda = if (indexVariable != null)
+            generateLambda(mapping, indexVariable, inputVariable)
+        else
+            generateLambda(inputVariable, mapping)
         return chainedCallGenerator.generate("$functionName($0) $1:'{}'", targetCollection, lambda)
     }
 
@@ -318,17 +329,18 @@ class MapToTransformation private constructor(
         fun create(
                 loop: KtForExpression,
                 inputVariable: KtCallableDeclaration,
+                indexVariable: KtCallableDeclaration?,
                 targetCollection: KtExpression,
                 mapping: KtExpression,
                 mapNotNull: Boolean
         ): ResultTransformation {
             val initialization = targetCollection.detectInitializationBeforeLoop(loop, checkNoOtherUsagesInLoop = true)
             if (initialization != null && initialization.initializer.hasNoSideEffect()) {
-                val transformation = MapToTransformation(loop, inputVariable, initialization.initializer, mapping, mapNotNull)
+                val transformation = MapToTransformation(loop, inputVariable, indexVariable, initialization.initializer, mapping, mapNotNull)
                 return AssignToVariableResultTransformation.createDelegated(transformation, initialization)
             }
             else {
-                return MapToTransformation(loop, inputVariable, targetCollection, mapping, mapNotNull)
+                return MapToTransformation(loop, inputVariable, indexVariable, targetCollection, mapping, mapNotNull)
             }
         }
     }

idea/testData/intentions/loopToCallChain/indexWithNestedLoop.kt

@@ -1,5 +1,5 @@
 // WITH_RUNTIME
-// INTENTION_TEXT: "Replace with '+= ...filterNot{}.mapIndexed{}'"
+// INTENTION_TEXT: "Replace with 'flatMap{}.filterNot{}.mapIndexedTo(){}'"
 fun foo(list: List<String>, target: MutableCollection<Int>) {
     var i = 0
     <caret>for (s in list) {

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

@@ -1,8 +1,8 @@
 // WITH_RUNTIME
-// INTENTION_TEXT: "Replace with '+= ...filterNot{}.mapIndexed{}'"
+// INTENTION_TEXT: "Replace with 'flatMap{}.filterNot{}.mapIndexedTo(){}'"
 fun foo(list: List<String>, target: MutableCollection<Int>) {
-    <caret>target += list
+    <caret>list
             .flatMap { it.indices }
             .filterNot { it == 10 }
-            .mapIndexed { i, j -> i + j }
+            .mapIndexedTo(target) { i, j -> i + j }
 }

idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo.kt

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'mapIndexedNotNullTo(){}'"
+fun foo(list: List<String?>, target: MutableList<Int>) {
+    <caret>for ((index, s) in list.withIndex()) {
+        val length = s?.substring(index)?.length
+        if (length == null) continue
+        target.add(length)
+    }
+}

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

@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'mapIndexedNotNullTo(){}'"
+fun foo(list: List<String?>, target: MutableList<Int>) {
+    <caret>list.mapIndexedNotNullTo(target) { index, s -> s?.substring(index)?.length }
+}

idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo_elvisContinue.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'mapIndexedNotNullTo(){}'"
+fun foo(list: List<String?>, target: MutableList<Int>) {
+    <caret>for ((index, s) in list.withIndex()) {
+        val length = s?.substring(index)?.length ?: continue
+        target.add(length)
+    }
+}

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

@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'mapIndexedNotNullTo(){}'"
+fun foo(list: List<String?>, target: MutableList<Int>) {
+    <caret>list.mapIndexedNotNullTo(target) { index, s -> s?.substring(index)?.length }
+}

idea/testData/intentions/loopToCallChain/mapIndexedTo.kt

@@ -1,5 +1,5 @@
 // WITH_RUNTIME
-// INTENTION_TEXT: "Replace with '+= mapIndexed{}'"
+// INTENTION_TEXT: "Replace with 'mapIndexedTo(){}'"
 fun foo(list: List<String>, target: MutableList<Int>) {
     <caret>for ((index, s) in list.withIndex()) {
         target.add(s.hashCode() * index)

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

@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'mapIndexedTo(){}'"
+fun foo(list: List<String>, target: MutableList<Int>) {
+    <caret>list.mapIndexedTo(target) { index, s -> s.hashCode() * index }
+}

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

@@ -1,5 +0,0 @@
-// WITH_RUNTIME
-// INTENTION_TEXT: "Replace with '+= mapIndexed{}'"
-fun foo(list: List<String>, target: MutableList<Int>) {
-    <caret>target += list.mapIndexed { index, s -> s.hashCode() * index }
-}