AddToCollectionTransformation checks that the result collection type cannot be changed

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

@@ -16,12 +16,20 @@
 
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain.result
 
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.VariableDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
+import org.jetbrains.kotlin.idea.imports.importableFqName
 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.idea.intentions.loopToCallChain.sequence.MapTransformation
+import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getCallNameExpression
+import org.jetbrains.kotlin.renderer.render
 
 class AddToCollectionTransformation(
         loop: KtForExpression,
@@ -69,42 +77,11 @@ class AddToCollectionTransformation(
             //TODO: check that it's MutableCollection's add
             val argument = callExpression.valueArguments.singleOrNull() ?: return null
             val argumentValue = argument.getArgumentExpression() ?: return null
-            val argumentIsInputVariable = argumentValue.isVariableReference(state.inputVariable)
 
-            //TODO: collection can be used as mutable collection or even ArrayList!
+            matchWithCollectionInitializationReplacement(state, targetCollection, argumentValue)
-            val collectionInitialization = targetCollection.detectInitializationBeforeLoop(state.outerLoop)
+                    ?.let { return it }
-            if (collectionInitialization != null) {
-                val collectionKind = collectionInitialization.initializer.isSimpleCollectionInstantiation()
-                when (collectionKind) {
-                    CollectionKind.LIST -> {
-                        val transformation = if (argumentIsInputVariable) {
-                            AssignToListTransformation(state.outerLoop, state.inputVariable, collectionInitialization)
-                        }
-                        else {
-                            val mapTransformation = MapTransformation(state.outerLoop, state.inputVariable, argumentValue)
-                            AssignSequenceTransformationResultTransformation(mapTransformation, collectionInitialization)
-                        }
-                        return ResultTransformationMatch(transformation)
-                    }
 
-                    CollectionKind.SET -> {
+            val transformation = if (argumentValue.isVariableReference(state.inputVariable)) {
-                        if (argumentIsInputVariable) {
-                            val transformation = AssignToSetTransformation(state.outerLoop, state.inputVariable, collectionInitialization)
-                            return ResultTransformationMatch(transformation)
-                        }
-                        else {
-                            val mapTransformation = MapTransformation(state.outerLoop, state.inputVariable, argumentValue)
-                            val transformation = AssignToSetTransformation(
-                                    state.outerLoop,
-                                    state.inputVariable/*TODO: it's not correct and it looks like not all transformations should have inputVariable*/,
-                                    collectionInitialization)
-                            return ResultTransformationMatch(transformation, mapTransformation)
-                        }
-                    }
-                }
-            }
-
-            val transformation = if (argumentIsInputVariable) {
                 AddToCollectionTransformation(state.outerLoop, state.inputVariable, targetCollection)
             }
             else {
@@ -112,6 +89,73 @@ class AddToCollectionTransformation(
             }
             return ResultTransformationMatch(transformation)
         }
+
+        private fun matchWithCollectionInitializationReplacement(
+                state: MatchingState,
+                targetCollection: KtExpression,
+                addOperationArgument: KtExpression
+        ): ResultTransformationMatch? {
+            val collectionInitialization = targetCollection.detectInitializationBeforeLoop(state.outerLoop) ?: return null
+            val collectionKind = collectionInitialization.initializer.isSimpleCollectionInstantiation() ?: return null
+            val argumentIsInputVariable = addOperationArgument.isVariableReference(state.inputVariable)
+            when (collectionKind) {
+                CollectionKind.LIST -> {
+                    when {
+                        canChangeInitializerType(collectionInitialization, KotlinBuiltIns.FQ_NAMES.list, state.outerLoop) -> {
+                            val transformation = if (argumentIsInputVariable) {
+                                AssignToListTransformation(state.outerLoop, state.inputVariable, collectionInitialization)
+                            }
+                            else {
+                                val mapTransformation = MapTransformation(state.outerLoop, state.inputVariable, addOperationArgument)
+                                AssignSequenceTransformationResultTransformation(mapTransformation, collectionInitialization)
+                            }
+                            return ResultTransformationMatch(transformation)
+                        }
+
+                        canChangeInitializerType(collectionInitialization, KotlinBuiltIns.FQ_NAMES.mutableList, state.outerLoop) -> {
+                            if (argumentIsInputVariable) {
+                                val transformation = AssignToMutableListTransformation(state.outerLoop, state.inputVariable, collectionInitialization)
+                                return ResultTransformationMatch(transformation)
+                            }
+                        }
+                    }
+                }
+
+                CollectionKind.SET -> {
+                    val assignToSetTransformation = when {
+                        canChangeInitializerType(collectionInitialization, KotlinBuiltIns.FQ_NAMES.set, state.outerLoop) -> {
+                            AssignToSetTransformation(state.outerLoop, state.inputVariable/*TODO: it's not correct and it looks like not all transformations should have inputVariable*/, collectionInitialization)
+                        }
+
+                        canChangeInitializerType(collectionInitialization, KotlinBuiltIns.FQ_NAMES.mutableSet, state.outerLoop) -> {
+                            AssignToMutableSetTransformation(state.outerLoop, state.inputVariable, collectionInitialization)
+                        }
+
+                        else -> return null
+                    }
+
+                    if (argumentIsInputVariable) {
+                        return ResultTransformationMatch(assignToSetTransformation)
+                    }
+                    else {
+                        val mapTransformation = MapTransformation(state.outerLoop, state.inputVariable, addOperationArgument)
+                        return ResultTransformationMatch(assignToSetTransformation, mapTransformation)
+                    }
+                }
+            }
+
+            return null
+        }
+
+        private fun canChangeInitializerType(initialization: VariableInitialization, newTypeFqName: FqName, loop: KtForExpression): Boolean {
+            val currentType = (initialization.variable.resolveToDescriptor() as VariableDescriptor).type
+            if ((currentType.constructor.declarationDescriptor as? ClassDescriptor)?.importableFqName == newTypeFqName) return true // already of the required type
+
+            if (initialization.initializationStatement != initialization.variable) return false
+
+            val newTypeText = newTypeFqName.render() + IdeDescriptorRenderers.SOURCE_CODE.renderTypeArguments(currentType.arguments)
+            return canChangeLocalVariableType(initialization.variable, newTypeText, loop)
+        }
     }
 }
 
@@ -170,6 +214,17 @@ class AssignToListTransformation(
     }
 }
 
+class AssignToMutableListTransformation(
+        loop: KtForExpression,
+        inputVariable: KtCallableDeclaration,
+        initialization: VariableInitialization
+) : AssignToVariableResultTransformation(loop, inputVariable, initialization) {
+
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
+        return chainedCallGenerator.generate("toMutableList()")
+    }
+}
+
 class AssignToSetTransformation(
         loop: KtForExpression,
         inputVariable: KtCallableDeclaration,
@@ -180,3 +235,14 @@ class AssignToSetTransformation(
         return chainedCallGenerator.generate("toSet()")
     }
 }
+
+class AssignToMutableSetTransformation(
+        loop: KtForExpression,
+        inputVariable: KtCallableDeclaration,
+        initialization: VariableInitialization
+) : AssignToVariableResultTransformation(loop, inputVariable, initialization) {
+
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
+        return chainedCallGenerator.generate("toMutableSet()")
+    }
+}

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

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain
 
+import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.UserDataHolderBase
 import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.search.LocalSearchScope
@@ -23,10 +24,12 @@ import com.intellij.psi.search.searches.ReferencesSearch
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
+import org.jetbrains.kotlin.diagnostics.Severity
 import org.jetbrains.kotlin.idea.analysis.analyzeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
+import org.jetbrains.kotlin.idea.core.copied
 import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.imports.importableFqName
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.isNullExpression
@@ -38,6 +41,8 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.*
 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.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.typeUtil.TypeNullability
@@ -330,3 +335,43 @@ fun KtExpression.isSimpleCollectionInstantiation(): CollectionKind? {
         else -> null
     }
 }
+
+fun canChangeLocalVariableType(variable: KtProperty, newTypeText: String, loop: KtForExpression): Boolean {
+    val bindingContext = variable.analyze(BodyResolveMode.FULL)
+
+    // analyze the closest block which is not used as expression
+    val block = variable.parents
+                        .filterIsInstance<KtBlockExpression>()
+                        .firstOrNull { bindingContext[BindingContext.USED_AS_EXPRESSION, it] != true }
+                ?: return false
+
+    val KEY = Key<Unit>("KEY")
+    block.putCopyableUserData(KEY, Unit)
+    variable.putCopyableUserData(KEY, Unit)
+    loop.putCopyableUserData(KEY, Unit)
+
+    val fileCopy = block.containingFile.copied()
+    val blockCopy: KtBlockExpression
+    val variableCopy: KtProperty
+    val loopCopy: KtForExpression
+    try {
+        blockCopy = fileCopy.findDescendantOfType<KtBlockExpression> { it.getCopyableUserData(KEY) != null }!!
+        variableCopy = blockCopy.findDescendantOfType<KtProperty> { it.getCopyableUserData(KEY) != null }!!
+        loopCopy = blockCopy.findDescendantOfType<KtForExpression> { it.getCopyableUserData(KEY) != null }!!
+    }
+    finally {
+        block.putCopyableUserData(KEY, null)
+        variable.putCopyableUserData(KEY, null)
+        loop.putCopyableUserData(KEY, null)
+    }
+
+    variableCopy.typeReference = KtPsiFactory(block).createType(newTypeText)
+
+    val resolutionScope = block.getResolutionScope(bindingContext, block.getResolutionFacade())
+    val newBindingContext = blockCopy.analyzeInContext(scope = resolutionScope,
+                                                       contextExpression = block,
+                                                       dataFlowInfo = bindingContext.getDataFlowInfo(block),
+                                                       trace = DelegatingBindingTrace(bindingContext, "Temporary trace"))
+    //TODO: what if there were errors before?
+    return newBindingContext.diagnostics.none { it.severity == Severity.ERROR && !loopCopy.isAncestor(it.psiElement) }
+}

idea/testData/intentions/loopToCallChain/assignFilter2.kt

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

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

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+import java.util.ArrayList
+
+fun foo(list: List<String>, p: Int): List<String> {
+    return if (p > 0) {
+        <caret>val result = list.filter { it.length > 0 }
+        result
+    }
+    else {
+        emptyList()
+    }
+}

idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt

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

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

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

idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired2.kt

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

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

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+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 }
+        result
+    }
+    else {
+        ArrayList()
+    }
+}

idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt

@@ -0,0 +1,14 @@
+// WITH_RUNTIME
+import java.util.*
+
+fun foo(list: List<String>): ArrayList<String> {
+    return run {
+        val result = ArrayList<String>()
+        <caret>for (s in list) {
+            if (s.length > 0) {
+                result.add(s)
+            }
+        }
+        result
+    }
+}

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

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

idea/testData/intentions/loopToCallChain/assignFilter_MutableListRequired.kt

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

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

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

idea/testData/intentions/loopToCallChain/toMutableSet.kt

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+import java.util.HashSet
+
+fun foo(map: Map<Int, String>): MutableCollection<String> {
+    val result = HashSet<String>()
+    <caret>for (s in map.values) {
+        result.add(s)
+    }
+    return result
+}

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

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+import java.util.HashSet
+
+fun foo(map: Map<Int, String>): MutableCollection<String> {
+    <caret>val result = map.values.toMutableSet()
+    return result
+}