diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt index c388b06d220..f43523676c8 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt @@ -30,6 +30,8 @@ interface Transformation { } interface SequenceTransformation : Transformation { + fun mergeWithPrevious(previousTransformation: SequenceTransformation): SequenceTransformation? = null + val affectsIndex: Boolean fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt index 35d74365c6a..207d4eb2305 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt @@ -71,8 +71,9 @@ fun match(loop: KtForExpression): ResultTransformationMatch? { val match = matcher.match(state) if (match != null) { sequenceTransformations.addAll(match.sequenceTransformations) - val resultMatch = ResultTransformationMatch(match.resultTransformation, sequenceTransformations) - return resultMatch.check { checkSmartCastsPreserved(loop, it) } + return ResultTransformationMatch(match.resultTransformation, sequenceTransformations) + .let { mergeTransformations(it) } + .check { checkSmartCastsPreserved(loop, it) } } } @@ -182,3 +183,29 @@ private fun ResultTransformationMatch.generateCallChain(loop: KtForExpression): return callChain } +private fun mergeTransformations(match: ResultTransformationMatch): ResultTransformationMatch { + val transformations = ArrayList().apply { addAll(match.sequenceTransformations); add(match.resultTransformation) } + + var anyChange: Boolean + do { + anyChange = false + for (index in 0..transformations.lastIndex - 1) { + val transformation = transformations[index] as SequenceTransformation + val next = transformations[index + 1] + val merged = when (next) { + is SequenceTransformation -> next.mergeWithPrevious(transformation) + is ResultTransformation -> next.mergeWithPrevious(transformation) + else -> error("Unknown transformation type: $next") + } ?: continue + + transformations[index] = merged + transformations.removeAt(index + 1) + anyChange = true + break + } + } while(anyChange) + + @Suppress("UNCHECKED_CAST") + return ResultTransformationMatch(transformations.last() as ResultTransformation, transformations.dropLast(1) as List) +} + diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndAssignTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndAssignTransformation.kt index f74a3cfd91c..6c2f1395142 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndAssignTransformation.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndAssignTransformation.kt @@ -34,7 +34,7 @@ class FindAndAssignTransformation( override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? { if (previousTransformation !is FilterTransformation) return null - if (filter != null) return null //TODO + assert(filter == null) { "Should not happen because no 2 consecutive FilterTransformation's possible"} return FindAndAssignTransformation(loop, previousTransformation.inputVariable, stdlibFunName, initialDeclaration, previousTransformation.buildRealCondition()) } diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndReturnTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndReturnTransformation.kt index 22e7a6cbbdd..f0a739bee37 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndReturnTransformation.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndReturnTransformation.kt @@ -34,7 +34,7 @@ class FindAndReturnTransformation( override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? { if (previousTransformation !is FilterTransformation) return null - if (filter != null) return null //TODO + assert(filter == null) { "Should not happen because no 2 consecutive FilterTransformation's possible"} return FindAndReturnTransformation(loop, previousTransformation.inputVariable, stdlibFunName, endReturn, previousTransformation.buildRealCondition()) } diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt index fe17dcc1ad9..3ef417ea542 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt @@ -29,12 +29,16 @@ class FilterTransformation( val isInverse: Boolean ) : SequenceTransformation { - init { - assert(condition.isPhysical) - } - fun buildRealCondition() = if (isInverse) condition.negate() else condition + override fun mergeWithPrevious(previousTransformation: SequenceTransformation): SequenceTransformation? { + if (previousTransformation !is FilterTransformation) return null + assert(previousTransformation.inputVariable == inputVariable) + val mergedCondition = KtPsiFactory(condition).createExpressionByPattern( + "$0 && $1", previousTransformation.buildRealCondition(), buildRealCondition()) + return FilterTransformation(inputVariable, mergedCondition, isInverse = false) //TODO: build filterNot in some cases? + } + override val affectsIndex: Boolean get() = true @@ -44,7 +48,6 @@ class FilterTransformation( return chainedCallGenerator.generate("$0$1:'{}'", name, lambda) } - //TODO: merge subsequent filters /** * Matches: * for (...) { diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FlatMapTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FlatMapTransformation.kt index 7d60b43c67c..07576db4c0c 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FlatMapTransformation.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FlatMapTransformation.kt @@ -30,10 +30,6 @@ class FlatMapTransformation( private val transform: KtExpression ) : SequenceTransformation { - init { - assert(transform.isPhysical) - } - override val affectsIndex: Boolean get() = true diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/MapTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/MapTransformation.kt index 5d03353b58f..1d1142852e8 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/MapTransformation.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/MapTransformation.kt @@ -26,10 +26,6 @@ class MapTransformation( val mapping: KtExpression ) : SequenceTransformation { - init { - assert(mapping.isPhysical) - } - override val affectsIndex: Boolean get() = false diff --git a/idea/testData/intentions/loopToCallChain/filter_mergeMultiple.kt b/idea/testData/intentions/loopToCallChain/filter_mergeMultiple.kt new file mode 100644 index 00000000000..0d400e62205 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/filter_mergeMultiple.kt @@ -0,0 +1,12 @@ +// WITH_RUNTIME +fun foo(list: List): String? { + for (s in list) { + if (s.isEmpty()) continue + if (s.length < 10 && s != "abc") { + if (s == "def") continue + val s1 = s + "x" + return s1 + } + } + return null +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/filter_mergeMultiple.kt.after b/idea/testData/intentions/loopToCallChain/filter_mergeMultiple.kt.after new file mode 100644 index 00000000000..3eda4d2a5b7 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/filter_mergeMultiple.kt.after @@ -0,0 +1,7 @@ +// WITH_RUNTIME +fun foo(list: List): String? { + return list + .filter { !it.isEmpty() && it.length < 10 && it != "abc" && it != "def" } + .map { it + "x" } + .firstOrNull() +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/firstOrNull_withMergedFilter.kt b/idea/testData/intentions/loopToCallChain/firstOrNull_withMergedFilter.kt new file mode 100644 index 00000000000..cdad088829d --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/firstOrNull_withMergedFilter.kt @@ -0,0 +1,11 @@ +// WITH_RUNTIME +fun foo(list: List): String? { + for (s in list) { + if (s.isEmpty()) continue + if (s.length < 10 && s != "abc") { + if (s == "def") continue + return s + } + } + return null +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/firstOrNull_withMergedFilter.kt.after b/idea/testData/intentions/loopToCallChain/firstOrNull_withMergedFilter.kt.after new file mode 100644 index 00000000000..1217edfac9b --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/firstOrNull_withMergedFilter.kt.after @@ -0,0 +1,4 @@ +// WITH_RUNTIME +fun foo(list: List): String? { + return list.firstOrNull { !it.isEmpty() && it.length < 10 && it != "abc" && it != "def" } +} \ No newline at end of file