Sometimes shorter code with filterTo, mapTo and flatMapTo

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/AddToCollectionTransformation.kt

@@ -40,15 +40,15 @@ class AddToCollectionTransformation(
     override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
         return when (previousTransformation) {
             is FilterTransformation -> {
-                FilterToTransformation(loop, inputVariable, targetCollection, previousTransformation.effectiveCondition()) //TODO: use filterNotTo?
+                FilterToTransformation.create(loop, inputVariable, targetCollection, previousTransformation.effectiveCondition()) //TODO: use filterNotTo?
             }
 
             is MapTransformation -> {
-                MapToTransformation(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.mapping)
+                MapToTransformation.create(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.mapping)
             }
 
             is FlatMapTransformation -> {
-                FlatMapToTransformation(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.transform)
+                FlatMapToTransformation.create(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.transform)
             }
 
             else -> null
@@ -89,7 +89,7 @@ class AddToCollectionTransformation(
                 AddToCollectionTransformation(state.outerLoop, state.inputVariable, targetCollection)
             }
             else {
-                MapToTransformation(state.outerLoop, state.inputVariable, targetCollection, argumentValue)
+                MapToTransformation.create(state.outerLoop, state.inputVariable, targetCollection, argumentValue)
             }
             return ResultTransformationMatch(transformation)
         }
@@ -163,7 +163,7 @@ class AddToCollectionTransformation(
     }
 }
 
-class FilterToTransformation(
+class FilterToTransformation private constructor(
         loop: KtForExpression,
         inputVariable: KtCallableDeclaration,
         private val targetCollection: KtExpression,
@@ -174,9 +174,39 @@ class FilterToTransformation(
         val lambda = generateLambda(inputVariable, filter)
         return chainedCallGenerator.generate("filterTo($0) $1:'{}'", targetCollection, lambda)
     }
+
+    companion object {
+        fun create(
+                loop: KtForExpression,
+                inputVariable: KtCallableDeclaration,
+                targetCollection: KtExpression,
+                filter: KtExpression
+        ): ResultTransformation {
+            val initialization = targetCollection.detectInitializationBeforeLoop(loop)
+            if (initialization != null && initialization.initializer.hasNoSideEffect()) {
+                return AssignFilterToTransformation(loop, inputVariable, initialization, filter)
+            }
+            else {
+                return FilterToTransformation(loop, inputVariable, targetCollection, filter)
+            }
+        }
+    }
 }
 
-class MapToTransformation(
+class AssignFilterToTransformation(
+        loop: KtForExpression,
+        inputVariable: KtCallableDeclaration,
+        targetCollectionInitialization: VariableInitialization,
+        private val filter: KtExpression
+) : AssignToVariableResultTransformation(loop, inputVariable, targetCollectionInitialization) {
+
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
+        val lambda = generateLambda(inputVariable, filter)
+        return chainedCallGenerator.generate("filterTo($0) $1:'{}'", initialization.initializer, lambda)
+    }
+}
+
+class MapToTransformation private constructor(
         loop: KtForExpression,
         inputVariable: KtCallableDeclaration,
         private val targetCollection: KtExpression,
@@ -187,9 +217,39 @@ class MapToTransformation(
         val lambda = generateLambda(inputVariable, mapping)
         return chainedCallGenerator.generate("mapTo($0) $1:'{}'", targetCollection, lambda)
     }
+
+    companion object {
+        fun create(
+                loop: KtForExpression,
+                inputVariable: KtCallableDeclaration,
+                targetCollection: KtExpression,
+                mapping: KtExpression
+        ): ResultTransformation {
+            val initialization = targetCollection.detectInitializationBeforeLoop(loop)
+            if (initialization != null && initialization.initializer.hasNoSideEffect()) {
+                return AssignMapToTransformation(loop, inputVariable, initialization, mapping)
+            }
+            else {
+                return MapToTransformation(loop, inputVariable, targetCollection, mapping)
+            }
+        }
+    }
 }
 
-class FlatMapToTransformation(
+class AssignMapToTransformation(
+        loop: KtForExpression,
+        inputVariable: KtCallableDeclaration,
+        targetCollectionInitialization: VariableInitialization,
+        private val mapping: KtExpression
+) : AssignToVariableResultTransformation(loop, inputVariable, targetCollectionInitialization) {
+
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
+        val lambda = generateLambda(inputVariable, mapping)
+        return chainedCallGenerator.generate("mapTo($0) $1:'{}'", initialization.initializer, lambda)
+    }
+}
+
+class FlatMapToTransformation private constructor(
         loop: KtForExpression,
         inputVariable: KtCallableDeclaration,
         private val targetCollection: KtExpression,
@@ -200,6 +260,36 @@ class FlatMapToTransformation(
         val lambda = generateLambda(inputVariable, transform)
         return chainedCallGenerator.generate("flatMapTo($0) $1:'{}'", targetCollection, lambda)
     }
+
+    companion object {
+        fun create(
+                loop: KtForExpression,
+                inputVariable: KtCallableDeclaration,
+                targetCollection: KtExpression,
+                transform: KtExpression
+        ): ResultTransformation {
+            val initialization = targetCollection.detectInitializationBeforeLoop(loop)
+            if (initialization != null && initialization.initializer.hasNoSideEffect()) {
+                return AssignFlatMapToTransformation(loop, inputVariable, initialization, transform)
+            }
+            else {
+                return FlatMapToTransformation(loop, inputVariable, targetCollection, transform)
+            }
+        }
+    }
+}
+
+class AssignFlatMapToTransformation(
+        loop: KtForExpression,
+        inputVariable: KtCallableDeclaration,
+        targetCollectionInitialization: VariableInitialization,
+        private val transform: KtExpression
+) : AssignToVariableResultTransformation(loop, inputVariable, targetCollectionInitialization) {
+
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
+        val lambda = generateLambda(inputVariable, transform)
+        return chainedCallGenerator.generate("flatMapTo($0) $1:'{}'", initialization.initializer, lambda)
+    }
 }
 
 class AssignToListTransformation(

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt

@@ -44,6 +44,7 @@ import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DelegatingBindingTrace
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.typeUtil.TypeNullability
 import org.jetbrains.kotlin.types.typeUtil.nullability
@@ -375,3 +376,26 @@ fun canChangeLocalVariableType(variable: KtProperty, newTypeText: String, loop:
     //TODO: what if there were errors before?
     return newBindingContext.diagnostics.none { it.severity == Severity.ERROR && !loopCopy.isAncestor(it.psiElement) }
 }
+
+private val NO_SIDE_EFFECT_STANDARD_CLASSES = setOf(
+        "java.util.ArrayList",
+        "java.util.LinkedList",
+        "java.util.HashSet",
+        "java.util.LinkedHashSet",
+        "java.util.HashMap",
+        "java.util.LinkedHashMap"
+)
+
+fun KtExpression.hasNoSideEffect(): Boolean {
+    val bindingContext = analyze(BodyResolveMode.PARTIAL)
+    if (ConstantExpressionEvaluator.getConstant(this, bindingContext) != null) return true
+
+    val callExpression = this as? KtCallExpression ?: return false//TODO: it can be qualified too
+    if (callExpression.valueArguments.any { it.getArgumentExpression()?.hasNoSideEffect() == false }) return false
+
+    val resolvedCall = callExpression.getResolvedCall(bindingContext) ?: return false
+    val constructorDescriptor = resolvedCall.resultingDescriptor as? ConstructorDescriptor ?: return false
+    val classDescriptor = constructorDescriptor.containingDeclaration
+    val classFqName = classDescriptor.importableFqName?.asString()
+    return classFqName in NO_SIDE_EFFECT_STANDARD_CLASSES
+}

idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt.after

@@ -2,7 +2,6 @@
 import java.util.ArrayList
 
 fun foo(list: List<String>): ArrayList<String> {
-    val result = ArrayList<String>()
-    <caret>list.filterTo(result) { it.length > 0 }
+    <caret>val result = list.filterTo(ArrayList<String>()) { it.length > 0 }
     return result
 }

idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired2.kt.after

@@ -3,8 +3,7 @@ import java.util.ArrayList
 
 fun foo(list: List<String>, p: Int): ArrayList<String> {
     return if (p > 0) {
-        val result = ArrayList<String>()
-        <caret>list.filterTo(result) { it.length > 0 }
+        <caret>val result = list.filterTo(ArrayList<String>()) { it.length > 0 }
         result
     }
     else {

idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt.after

@@ -3,8 +3,7 @@ import java.util.*
 
 fun foo(list: List<String>): ArrayList<String> {
     return run {
-        val result = ArrayList<String>()
-        <caret>list.filterTo(result) { it.length > 0 }
+        <caret>val result = list.filterTo(ArrayList<String>()) { it.length > 0 }
         result
     }
 }

idea/testData/intentions/loopToCallChain/filterTo2.kt

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+fun foo(list: List<String>) {
+    val target = createCollection()
+    <caret>for (s in list) {
+        if (s.length > 0)
+            target.add(s)
+    }
+}
+
+fun createCollection() = java.util.ArrayList<String>()

idea/testData/intentions/loopToCallChain/filterTo2.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+fun foo(list: List<String>) {
+    val target = createCollection()
+    <caret>list.filterTo(target) { it.length > 0 }
+}
+
+fun createCollection() = java.util.ArrayList<String>()

idea/testData/intentions/loopToCallChain/flatMapTo2.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+import java.util.ArrayList
+
+fun foo(list: List<String>): List<String> {
+    val target = ArrayList<String>(100)
+    <caret>for (s in list) {
+        for (line in s.lines()) {
+            target.add(line)
+        }
+    }
+    return target
+}

idea/testData/intentions/loopToCallChain/flatMapTo2.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+import java.util.ArrayList
+
+fun foo(list: List<String>): List<String> {
+    <caret>val target = list.flatMapTo(ArrayList<String>(100)) { it.lines() }
+    return target
+}

idea/testData/intentions/loopToCallChain/flatMapTo3.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+fun foo(list: List<String>): List<String> {
+    val target = createCollection()
+    <caret>for (s in list) {
+        for (line in s.lines()) {
+            target.add(line)
+        }
+    }
+    return target
+}
+
+fun createCollection() = java.util.ArrayList<String>()

idea/testData/intentions/loopToCallChain/flatMapTo3.kt.after

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun foo(list: List<String>): List<String> {
+    val target = createCollection()
+    <caret>list.flatMapTo(target) { it.lines() }
+    return target
+}
+
+fun createCollection() = java.util.ArrayList<String>()

idea/testData/intentions/loopToCallChain/mapTo3.kt

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+import java.util.ArrayList
+
+fun foo(list: List<String>): List<Int> {
+    val target = ArrayList<Int>(100)
+    <caret>for (s in list) {
+        if (s.length > 0)
+            target.add(s.hashCode())
+    }
+    return target
+}

idea/testData/intentions/loopToCallChain/mapTo3.kt.after

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+import java.util.ArrayList
+
+fun foo(list: List<String>): List<Int> {
+    <caret>val target = list
+            .filter { it.length > 0 }
+            .mapTo(ArrayList<Int>(100)) { it.hashCode() }
+    return target
+}

idea/testData/intentions/loopToCallChain/mapTo4.kt

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+fun foo(list: List<String>): List<Int> {
+    val target = createCollection()
+    <caret>for (s in list) {
+        if (s.length > 0)
+            target.add(s.hashCode())
+    }
+    return target
+}
+
+fun createCollection(): MutableList<Int> = java.util.ArrayList()

idea/testData/intentions/loopToCallChain/mapTo4.kt.after

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+fun foo(list: List<String>): List<Int> {
+    val target = createCollection()
+    <caret>list
+            .filter { it.length > 0 }
+            .mapTo(target) { it.hashCode() }
+    return target
+}
+
+fun createCollection(): MutableList<Int> = java.util.ArrayList()