flatMapTo supported

2016-04-06 17:23:17 +03:00
parent f3c760607b
commit d3721e9462
4 changed files with 38 additions and 3 deletions
@@ -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)
+    }
 }
@@ -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
@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo(list: List<String>, target: MutableList<String>) {
+    <caret>for (s in list) {
+        for (line in s.lines()) {
+            target.add(line)
+        }
+    }
+}
@@ -0,0 +1,4 @@
+// WITH_RUNTIME
+fun foo(list: List<String>, target: MutableList<String>) {
+    <caret>list.flatMapTo(target) { it.lines() }
+}