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 23bdb827302..ecb60783dad 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,7 +49,11 @@ class AddToCollectionTransformation( } is MapTransformation -> { - MapToTransformation.create(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.mapping) + MapToTransformation.create(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.mapping, mapNotNull = false) + } + + is MapNotNullTransformation -> { + MapToTransformation.create(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.mapping, mapNotNull = true) } is FlatMapTransformation -> { @@ -116,7 +120,7 @@ class AddToCollectionTransformation( return ResultTransformationMatch(addToCollectionTransformation, mapIndexedTransformation) } else { - return ResultTransformationMatch(MapToTransformation.create(state.outerLoop, state.inputVariable, targetCollection, argumentValue)) + return ResultTransformationMatch(MapToTransformation.create(state.outerLoop, state.inputVariable, targetCollection, argumentValue, mapNotNull = false)) } } @@ -296,15 +300,18 @@ class MapToTransformation private constructor( loop: KtForExpression, private val inputVariable: KtCallableDeclaration, private val targetCollection: KtExpression, - private val mapping: KtExpression + private val mapping: KtExpression, + mapNotNull: Boolean ) : ReplaceLoopResultTransformation(loop) { + private val functionName = if (mapNotNull) "mapNotNullTo" else "mapTo" + override val presentation: String - get() = "mapTo(){}" + get() = "$functionName(){}" override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression { val lambda = generateLambda(inputVariable, mapping) - return chainedCallGenerator.generate("mapTo($0) $1:'{}'", targetCollection, lambda) + return chainedCallGenerator.generate("$functionName($0) $1:'{}'", targetCollection, lambda) } companion object { @@ -312,15 +319,16 @@ class MapToTransformation private constructor( loop: KtForExpression, inputVariable: KtCallableDeclaration, targetCollection: KtExpression, - mapping: 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) + val transformation = MapToTransformation(loop, inputVariable, initialization.initializer, mapping, mapNotNull) return AssignToVariableResultTransformation.createDelegated(transformation, initialization) } else { - return MapToTransformation(loop, inputVariable, targetCollection, mapping) + return MapToTransformation(loop, inputVariable, targetCollection, mapping, mapNotNull) } } } diff --git a/idea/testData/intentions/loopToCallChain/mapNotNullTo.kt b/idea/testData/intentions/loopToCallChain/mapNotNullTo.kt new file mode 100644 index 00000000000..817042c1f0e --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapNotNullTo.kt @@ -0,0 +1,9 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapNotNullTo(){}'" +fun foo(list: List, target: MutableList) { + for (s in list) { + val length = s?.length + if (length == null) continue + target.add(length) + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapNotNullTo.kt.after b/idea/testData/intentions/loopToCallChain/mapNotNullTo.kt.after new file mode 100644 index 00000000000..cc1630fffef --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapNotNullTo.kt.after @@ -0,0 +1,5 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapNotNullTo(){}'" +fun foo(list: List, target: MutableList) { + list.mapNotNullTo(target) { it?.length } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapNotNullTo_elvisContinue.kt b/idea/testData/intentions/loopToCallChain/mapNotNullTo_elvisContinue.kt new file mode 100644 index 00000000000..ad7d20cfdec --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapNotNullTo_elvisContinue.kt @@ -0,0 +1,8 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapNotNullTo(){}'" +fun foo(list: List, target: MutableList) { + for (s in list) { + val length = s?.length ?: continue + target.add(length) + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapNotNullTo_elvisContinue.kt.after b/idea/testData/intentions/loopToCallChain/mapNotNullTo_elvisContinue.kt.after new file mode 100644 index 00000000000..cc1630fffef --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapNotNullTo_elvisContinue.kt.after @@ -0,0 +1,5 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapNotNullTo(){}'" +fun foo(list: List, target: MutableList) { + list.mapNotNullTo(target) { it?.length } +} \ No newline at end of file