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 b825ba8543a..d34ed91094f 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 @@ -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! - val collectionInitialization = targetCollection.detectInitializationBeforeLoop(state.outerLoop) - 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) - } + matchWithCollectionInitializationReplacement(state, targetCollection, argumentValue) + ?.let { return it } - CollectionKind.SET -> { - 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) { + val transformation = if (argumentValue.isVariableReference(state.inputVariable)) { 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()") + } +} diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt index 5ca1561ac90..854d970243e 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt +++ b/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() + .firstOrNull { bindingContext[BindingContext.USED_AS_EXPRESSION, it] != true } + ?: return false + + val KEY = Key("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 { it.getCopyableUserData(KEY) != null }!! + variableCopy = blockCopy.findDescendantOfType { it.getCopyableUserData(KEY) != null }!! + loopCopy = blockCopy.findDescendantOfType { 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) } +} diff --git a/idea/testData/intentions/loopToCallChain/assignFilter2.kt b/idea/testData/intentions/loopToCallChain/assignFilter2.kt new file mode 100644 index 00000000000..f756a363cdb --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/assignFilter2.kt @@ -0,0 +1,17 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List, p: Int): List { + return if (p > 0) { + val result = ArrayList() + for (s in list) { + if (s.length > 0) { + result.add(s) + } + } + result + } + else { + emptyList() + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/assignFilter2.kt.after b/idea/testData/intentions/loopToCallChain/assignFilter2.kt.after new file mode 100644 index 00000000000..a20e01608db --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/assignFilter2.kt.after @@ -0,0 +1,12 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List, p: Int): List { + return if (p > 0) { + val result = list.filter { it.length > 0 } + result + } + else { + emptyList() + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt new file mode 100644 index 00000000000..1366a23a335 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt @@ -0,0 +1,12 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List): ArrayList { + val result = ArrayList() + for (s in list) { + if (s.length > 0) { + result.add(s) + } + } + return result +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt.after b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt.after new file mode 100644 index 00000000000..d45e17f1fd4 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt.after @@ -0,0 +1,8 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List): ArrayList { + val result = ArrayList() + list.filterTo(result) { it.length > 0 } + return result +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired2.kt b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired2.kt new file mode 100644 index 00000000000..890d4f8456b --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired2.kt @@ -0,0 +1,17 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List, p: Int): ArrayList { + return if (p > 0) { + val result = ArrayList() + for (s in list) { + if (s.length > 0) { + result.add(s) + } + } + result + } + else { + ArrayList() + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired2.kt.after b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired2.kt.after new file mode 100644 index 00000000000..143092f7312 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired2.kt.after @@ -0,0 +1,13 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List, p: Int): ArrayList { + return if (p > 0) { + val result = ArrayList() + list.filterTo(result) { it.length > 0 } + result + } + else { + ArrayList() + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt new file mode 100644 index 00000000000..393ca2b4d2e --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt @@ -0,0 +1,14 @@ +// WITH_RUNTIME +import java.util.* + +fun foo(list: List): ArrayList { + return run { + val result = ArrayList() + for (s in list) { + if (s.length > 0) { + result.add(s) + } + } + result + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt.after b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt.after new file mode 100644 index 00000000000..77be9be6c81 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt.after @@ -0,0 +1,10 @@ +// WITH_RUNTIME +import java.util.* + +fun foo(list: List): ArrayList { + return run { + val result = ArrayList() + list.filterTo(result) { it.length > 0 } + result + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/assignFilter_MutableListRequired.kt b/idea/testData/intentions/loopToCallChain/assignFilter_MutableListRequired.kt new file mode 100644 index 00000000000..e2458bd74bc --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/assignFilter_MutableListRequired.kt @@ -0,0 +1,12 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List): MutableList { + val result = ArrayList() + for (s in list) { + if (s.length > 0) { + result.add(s) + } + } + return result +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/assignFilter_MutableListRequired.kt.after b/idea/testData/intentions/loopToCallChain/assignFilter_MutableListRequired.kt.after new file mode 100644 index 00000000000..b7117e200b3 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/assignFilter_MutableListRequired.kt.after @@ -0,0 +1,9 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List): MutableList { + val result = list + .filter { it.length > 0 } + .toMutableList() + return result +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/toMutableSet.kt b/idea/testData/intentions/loopToCallChain/toMutableSet.kt new file mode 100644 index 00000000000..3a1446efa1d --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/toMutableSet.kt @@ -0,0 +1,10 @@ +// WITH_RUNTIME +import java.util.HashSet + +fun foo(map: Map): MutableCollection { + val result = HashSet() + for (s in map.values) { + result.add(s) + } + return result +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/toMutableSet.kt.after b/idea/testData/intentions/loopToCallChain/toMutableSet.kt.after new file mode 100644 index 00000000000..2b5241c540d --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/toMutableSet.kt.after @@ -0,0 +1,7 @@ +// WITH_RUNTIME +import java.util.HashSet + +fun foo(map: Map): MutableCollection { + val result = map.values.toMutableSet() + return result +} \ No newline at end of file