diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/AddToCollectionTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/AddToCollectionTransformation.kt index ecb60783dad..31d6677de35 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/AddToCollectionTransformation.kt +++ b/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) } } } diff --git a/idea/testData/intentions/loopToCallChain/indexWithNestedLoop.kt b/idea/testData/intentions/loopToCallChain/indexWithNestedLoop.kt index 60a30e93b0b..9e2b60d1ba2 100644 --- a/idea/testData/intentions/loopToCallChain/indexWithNestedLoop.kt +++ b/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, target: MutableCollection) { var i = 0 for (s in list) { diff --git a/idea/testData/intentions/loopToCallChain/indexWithNestedLoop.kt.after b/idea/testData/intentions/loopToCallChain/indexWithNestedLoop.kt.after index 16d2886712d..dd5158f4346 100644 --- a/idea/testData/intentions/loopToCallChain/indexWithNestedLoop.kt.after +++ b/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, target: MutableCollection) { - target += list + list .flatMap { it.indices } .filterNot { it == 10 } - .mapIndexed { i, j -> i + j } + .mapIndexedTo(target) { i, j -> i + j } } \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo.kt b/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo.kt new file mode 100644 index 00000000000..f2da84b4806 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo.kt @@ -0,0 +1,9 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapIndexedNotNullTo(){}'" +fun foo(list: List, target: MutableList) { + for ((index, s) in list.withIndex()) { + val length = s?.substring(index)?.length + if (length == null) continue + target.add(length) + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo.kt.after b/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo.kt.after new file mode 100644 index 00000000000..489199c5f76 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo.kt.after @@ -0,0 +1,5 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapIndexedNotNullTo(){}'" +fun foo(list: List, target: MutableList) { + list.mapIndexedNotNullTo(target) { index, s -> s?.substring(index)?.length } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo_elvisContinue.kt b/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo_elvisContinue.kt new file mode 100644 index 00000000000..8e0006a4cab --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo_elvisContinue.kt @@ -0,0 +1,8 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapIndexedNotNullTo(){}'" +fun foo(list: List, target: MutableList) { + for ((index, s) in list.withIndex()) { + val length = s?.substring(index)?.length ?: continue + target.add(length) + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo_elvisContinue.kt.after b/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo_elvisContinue.kt.after new file mode 100644 index 00000000000..489199c5f76 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapIndexedNotNullTo_elvisContinue.kt.after @@ -0,0 +1,5 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapIndexedNotNullTo(){}'" +fun foo(list: List, target: MutableList) { + list.mapIndexedNotNullTo(target) { index, s -> s?.substring(index)?.length } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt b/idea/testData/intentions/loopToCallChain/mapIndexedTo.kt similarity index 76% rename from idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt rename to idea/testData/intentions/loopToCallChain/mapIndexedTo.kt index 6bcad90c726..ce7a4faa755 100644 --- a/idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt +++ b/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, target: MutableList) { for ((index, s) in list.withIndex()) { target.add(s.hashCode() * index) diff --git a/idea/testData/intentions/loopToCallChain/mapIndexedTo.kt.after b/idea/testData/intentions/loopToCallChain/mapIndexedTo.kt.after new file mode 100644 index 00000000000..6c96861dcf2 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapIndexedTo.kt.after @@ -0,0 +1,5 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapIndexedTo(){}'" +fun foo(list: List, target: MutableList) { + list.mapIndexedTo(target) { index, s -> s.hashCode() * index } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt.after b/idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt.after deleted file mode 100644 index fb139ec672e..00000000000 --- a/idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt.after +++ /dev/null @@ -1,5 +0,0 @@ -// WITH_RUNTIME -// INTENTION_TEXT: "Replace with '+= mapIndexed{}'" -fun foo(list: List, target: MutableList) { - target += list.mapIndexed { index, s -> s.hashCode() * index } -} \ No newline at end of file