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 3653a91662c..ae7ba0ba203 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 @@ -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( 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 854d970243e..a1312c73d39 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt +++ b/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 +} diff --git a/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt.after b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt.after index d45e17f1fd4..3c46f7c8610 100644 --- a/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt.after +++ b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired.kt.after @@ -2,7 +2,6 @@ import java.util.ArrayList fun foo(list: List): ArrayList { - val result = ArrayList() - list.filterTo(result) { it.length > 0 } + val result = list.filterTo(ArrayList()) { it.length > 0 } return result } \ 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 index 143092f7312..3a734969cfb 100644 --- a/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired2.kt.after +++ b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired2.kt.after @@ -3,8 +3,7 @@ import java.util.ArrayList fun foo(list: List, p: Int): ArrayList { return if (p > 0) { - val result = ArrayList() - list.filterTo(result) { it.length > 0 } + val result = list.filterTo(ArrayList()) { it.length > 0 } result } else { diff --git a/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt.after b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt.after index 77be9be6c81..2f78ebc5387 100644 --- a/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt.after +++ b/idea/testData/intentions/loopToCallChain/assignFilter_ArrayListRequired3.kt.after @@ -3,8 +3,7 @@ import java.util.* fun foo(list: List): ArrayList { return run { - val result = ArrayList() - list.filterTo(result) { it.length > 0 } + val result = list.filterTo(ArrayList()) { it.length > 0 } result } } \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/filterTo2.kt b/idea/testData/intentions/loopToCallChain/filterTo2.kt new file mode 100644 index 00000000000..0e4c50a82a1 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/filterTo2.kt @@ -0,0 +1,10 @@ +// WITH_RUNTIME +fun foo(list: List) { + val target = createCollection() + for (s in list) { + if (s.length > 0) + target.add(s) + } +} + +fun createCollection() = java.util.ArrayList() \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/filterTo2.kt.after b/idea/testData/intentions/loopToCallChain/filterTo2.kt.after new file mode 100644 index 00000000000..41a98d1c40c --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/filterTo2.kt.after @@ -0,0 +1,7 @@ +// WITH_RUNTIME +fun foo(list: List) { + val target = createCollection() + list.filterTo(target) { it.length > 0 } +} + +fun createCollection() = java.util.ArrayList() \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/flatMapTo2.kt b/idea/testData/intentions/loopToCallChain/flatMapTo2.kt new file mode 100644 index 00000000000..3536bc0ebad --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/flatMapTo2.kt @@ -0,0 +1,12 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List): List { + val target = ArrayList(100) + for (s in list) { + for (line in s.lines()) { + target.add(line) + } + } + return target +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/flatMapTo2.kt.after b/idea/testData/intentions/loopToCallChain/flatMapTo2.kt.after new file mode 100644 index 00000000000..b56dc54682d --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/flatMapTo2.kt.after @@ -0,0 +1,7 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List): List { + val target = list.flatMapTo(ArrayList(100)) { it.lines() } + return target +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/flatMapTo3.kt b/idea/testData/intentions/loopToCallChain/flatMapTo3.kt new file mode 100644 index 00000000000..82c0296d7a8 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/flatMapTo3.kt @@ -0,0 +1,12 @@ +// WITH_RUNTIME +fun foo(list: List): List { + val target = createCollection() + for (s in list) { + for (line in s.lines()) { + target.add(line) + } + } + return target +} + +fun createCollection() = java.util.ArrayList() \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/flatMapTo3.kt.after b/idea/testData/intentions/loopToCallChain/flatMapTo3.kt.after new file mode 100644 index 00000000000..6d8693689b4 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/flatMapTo3.kt.after @@ -0,0 +1,8 @@ +// WITH_RUNTIME +fun foo(list: List): List { + val target = createCollection() + list.flatMapTo(target) { it.lines() } + return target +} + +fun createCollection() = java.util.ArrayList() \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapTo3.kt b/idea/testData/intentions/loopToCallChain/mapTo3.kt new file mode 100644 index 00000000000..871617d12d0 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapTo3.kt @@ -0,0 +1,11 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List): List { + val target = ArrayList(100) + for (s in list) { + if (s.length > 0) + target.add(s.hashCode()) + } + return target +} diff --git a/idea/testData/intentions/loopToCallChain/mapTo3.kt.after b/idea/testData/intentions/loopToCallChain/mapTo3.kt.after new file mode 100644 index 00000000000..2d1ac0d20dd --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapTo3.kt.after @@ -0,0 +1,9 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(list: List): List { + val target = list + .filter { it.length > 0 } + .mapTo(ArrayList(100)) { it.hashCode() } + return target +} diff --git a/idea/testData/intentions/loopToCallChain/mapTo4.kt b/idea/testData/intentions/loopToCallChain/mapTo4.kt new file mode 100644 index 00000000000..9267ac1218f --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapTo4.kt @@ -0,0 +1,11 @@ +// WITH_RUNTIME +fun foo(list: List): List { + val target = createCollection() + for (s in list) { + if (s.length > 0) + target.add(s.hashCode()) + } + return target +} + +fun createCollection(): MutableList = java.util.ArrayList() \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapTo4.kt.after b/idea/testData/intentions/loopToCallChain/mapTo4.kt.after new file mode 100644 index 00000000000..2eeef470098 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapTo4.kt.after @@ -0,0 +1,10 @@ +// WITH_RUNTIME +fun foo(list: List): List { + val target = createCollection() + list + .filter { it.length > 0 } + .mapTo(target) { it.hashCode() } + return target +} + +fun createCollection(): MutableList = java.util.ArrayList() \ No newline at end of file