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 31d6677de35..4ee580ab2b6 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,19 +49,7 @@ class AddToCollectionTransformation( } is MapTransformation -> { - MapToTransformation.create(loop, previousTransformation.inputVariable, null, targetCollection, previousTransformation.mapping, mapNotNull = false) - } - - 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) + MapToTransformation.create(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, targetCollection, previousTransformation.mapping, previousTransformation.mapNotNull) } 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) } } diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt index 28b8048d51e..03354af1d86 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt +++ b/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) { - is MapTransformation -> MapNotNullTransformation(loop, previousTransformation.inputVariable, previousTransformation.mapping) - is MapIndexedTransformation -> MapIndexedNotNullTransformation(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, previousTransformation.mapping) - else -> null + if (previousTransformation is MapTransformation) { + return MapTransformation(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, previousTransformation.mapping, mapNotNull = true) } + return null } override val affectsIndex: Boolean diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FlatMapTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FlatMapTransformation.kt index 3ca2bb073d4..532bff08eb3 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FlatMapTransformation.kt +++ b/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( diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/MapTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/MapTransformation.kt index 25b1ef03804..a8748afe357 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/MapTransformation.kt +++ b/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) - return chainedCallGenerator.generate("map$0:'{}'", lambda) + val lambda = if (indexVariable != null) + 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) - } -}