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 e27521432c2..b224cc01df0 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 @@ -18,6 +18,7 @@ package org.jetbrains.kotlin.idea.intentions.loopToCallChain.result import org.jetbrains.kotlin.idea.intentions.loopToCallChain.* import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterTransformation +import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FlatMapTransformation import org.jetbrains.kotlin.psi.* import org.jetbrains.kotlin.psi.psiUtil.getCallNameExpression @@ -28,8 +29,17 @@ class AddToCollectionTransformation( ) : ReplaceLoopTransformation(loop, inputVariable) { override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? { - if (previousTransformation !is FilterTransformation) return null - return FilterToTransformation(loop, inputVariable, targetCollection, previousTransformation.effectiveCondition()) //TODO: use filterNotTo? + return when (previousTransformation) { + is FilterTransformation -> { + FilterToTransformation(loop, inputVariable, targetCollection, previousTransformation.effectiveCondition()) //TODO: use filterNotTo? + } + + is FlatMapTransformation -> { + FlatMapToTransformation(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.transform) + } + + else -> null + } } override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression { @@ -96,4 +106,17 @@ class MapToTransformation( val lambda = generateLambda(inputVariable, mapping) return chainedCallGenerator.generate("mapTo($0) $1:'{}'", targetCollection, lambda) } +} + +class FlatMapToTransformation( + loop: KtForExpression, + inputVariable: KtCallableDeclaration, + private val targetCollection: KtExpression, + private val transform: KtExpression +) : ReplaceLoopTransformation(loop, inputVariable) { + + override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression { + val lambda = generateLambda(inputVariable, transform) + return chainedCallGenerator.generate("flatMapTo($0) $1:'{}'", targetCollection, lambda) + } } \ No newline at end of file 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 fe47e5d1767..3a93c33d734 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 @@ -27,7 +27,7 @@ import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode class FlatMapTransformation( override val inputVariable: KtCallableDeclaration, - private val transform: KtExpression + val transform: KtExpression ) : SequenceTransformation { override val affectsIndex: Boolean diff --git a/idea/testData/intentions/loopToCallChain/flatMapTo.kt b/idea/testData/intentions/loopToCallChain/flatMapTo.kt new file mode 100644 index 00000000000..ff889d57ad6 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/flatMapTo.kt @@ -0,0 +1,8 @@ +// WITH_RUNTIME +fun foo(list: List, target: MutableList) { + for (s in list) { + for (line in s.lines()) { + target.add(line) + } + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/flatMapTo.kt.after b/idea/testData/intentions/loopToCallChain/flatMapTo.kt.after new file mode 100644 index 00000000000..cd41c91bafe --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/flatMapTo.kt.after @@ -0,0 +1,4 @@ +// WITH_RUNTIME +fun foo(list: List, target: MutableList) { + list.flatMapTo(target) { it.lines() } +} \ No newline at end of file