From 14e87b1f2c3e2f50bc405cf675b06e6132a098fe Mon Sep 17 00:00:00 2001 From: Valentin Kipyatkov Date: Wed, 20 Apr 2016 20:00:41 +0300 Subject: [PATCH] Initial support for indexed transformations --- .../LoopToCallChainIntention.kt | 2 +- .../intentions/loopToCallChain/interfaces.kt | 5 ++ .../loopToCallChain/matchAndConvert.kt | 73 +++++++++++++++---- .../result/AddToCollectionTransformation.kt | 28 ++++--- .../result/CountTransformation.kt | 3 + .../result/FindAndAssignTransformation.kt | 6 +- .../result/FindAndReturnTransformation.kt | 6 +- .../sequence/FilterTransformation.kt | 36 +++++++-- .../sequence/FlatMapTransformation.kt | 21 ++++-- .../sequence/MapTransformation.kt | 28 ++++++- .../idea/intentions/loopToCallChain/utils.kt | 21 ++++++ .../loopToCallChain/any_indexNeeded.kt | 11 +++ .../loopToCallChain/any_indexNeeded.kt.after | 7 ++ .../loopToCallChain/filterIndexed.kt | 10 +++ .../loopToCallChain/filterIndexed.kt.after | 7 ++ .../loopToCallChain/flatMap_indexUsed.kt | 10 +++ .../flatMap_indexUsed.kt.after | 8 ++ .../loopToCallChain/mapAndFilterIndexed.kt | 11 +++ .../mapAndFilterIndexed.kt.after | 8 ++ .../intentions/loopToCallChain/mapIndexed.kt | 9 +++ .../loopToCallChain/mapIndexed.kt.after | 7 ++ .../loopToCallChain/mapIndexed_twice.kt | 10 +++ .../loopToCallChain/mapIndexed_twice.kt.after | 8 ++ .../loopToCallChain/mapTo_indexUsed.kt | 7 ++ .../loopToCallChain/mapTo_indexUsed.kt.after | 5 ++ .../mapUsesOldIndexAfterFilter.kt | 10 +++ 26 files changed, 311 insertions(+), 46 deletions(-) create mode 100644 idea/testData/intentions/loopToCallChain/any_indexNeeded.kt create mode 100644 idea/testData/intentions/loopToCallChain/any_indexNeeded.kt.after create mode 100644 idea/testData/intentions/loopToCallChain/filterIndexed.kt create mode 100644 idea/testData/intentions/loopToCallChain/filterIndexed.kt.after create mode 100644 idea/testData/intentions/loopToCallChain/flatMap_indexUsed.kt create mode 100644 idea/testData/intentions/loopToCallChain/flatMap_indexUsed.kt.after create mode 100644 idea/testData/intentions/loopToCallChain/mapAndFilterIndexed.kt create mode 100644 idea/testData/intentions/loopToCallChain/mapAndFilterIndexed.kt.after create mode 100644 idea/testData/intentions/loopToCallChain/mapIndexed.kt create mode 100644 idea/testData/intentions/loopToCallChain/mapIndexed.kt.after create mode 100644 idea/testData/intentions/loopToCallChain/mapIndexed_twice.kt create mode 100644 idea/testData/intentions/loopToCallChain/mapIndexed_twice.kt.after create mode 100644 idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt create mode 100644 idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt.after create mode 100644 idea/testData/intentions/loopToCallChain/mapUsesOldIndexAfterFilter.kt diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/LoopToCallChainIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/LoopToCallChainIntention.kt index e981862aca5..f39a0a111d1 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/LoopToCallChainIntention.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/LoopToCallChainIntention.kt @@ -35,7 +35,7 @@ class LoopToCallChainIntention : SelfTargetingRangeIntention( ) { override fun applicabilityRange(element: KtForExpression): TextRange? { val match = match(element) - text = if (match != null) "Replace with '${match.buildPresentation()}'" else defaultText + text = if (match != null) "Replace with '${match.transformationMatch.buildPresentation()}'" else defaultText return if (match != null) element.forKeyword.textRange else null } 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 b40941e91c7..eb7a61b1192 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt @@ -86,6 +86,9 @@ data class MatchingState( val innerLoop: KtForExpression, val statements: Collection, val inputVariable: KtCallableDeclaration, + /** + * Matchers can assume that indexVariable is null if it's not used in the rest of the loop + */ val indexVariable: KtCallableDeclaration? ) @@ -113,6 +116,8 @@ class SequenceTransformationMatch( */ interface ResultTransformationMatcher { fun match(state: MatchingState): ResultTransformationMatch? + + val indexVariableUsePossible: Boolean } class ResultTransformationMatch( 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 ac3e9ad1cf3..623fa75d994 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt @@ -29,6 +29,7 @@ import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FlatMapTran import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.MapTransformation import org.jetbrains.kotlin.idea.util.CommentSaver import org.jetbrains.kotlin.idea.util.getResolutionScope +import org.jetbrains.kotlin.name.Name import org.jetbrains.kotlin.psi.* import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType @@ -54,14 +55,21 @@ object MatcherRegistrar { ) } -fun match(loop: KtForExpression): ResultTransformationMatch? { +data class MatchResult( + val sequenceExpression: KtExpression, + val transformationMatch: ResultTransformationMatch +) + +fun match(loop: KtForExpression): MatchResult? { + val (inputVariable, indexVariable, sequenceExpression) = extractLoopData(loop) ?: return null + val sequenceTransformations = ArrayList() var state = MatchingState( outerLoop = loop, innerLoop = loop, statements = listOf(loop.body ?: return null), - inputVariable = loop.loopParameter ?: return null, - indexVariable = null + inputVariable = inputVariable, + indexVariable = indexVariable ) MatchLoop@ @@ -73,7 +81,14 @@ fun match(loop: KtForExpression): ResultTransformationMatch? { val inputVariableUsed = state.inputVariable.hasUsages(state.statements) + // drop index variable if it's not used anymore + if (state.indexVariable != null && !state.indexVariable!!.hasUsages(state.statements)) { + state = state.copy(indexVariable = null) + } + for (matcher in MatcherRegistrar.resultMatchers) { + if (state.indexVariable != null && !matcher.indexVariableUsePossible) continue + val match = matcher.match(state) if (match != null) { if (!inputVariableUsed @@ -85,6 +100,7 @@ fun match(loop: KtForExpression): ResultTransformationMatch? { sequenceTransformations.addAll(match.sequenceTransformations) return ResultTransformationMatch(match.resultTransformation, sequenceTransformations) .let { mergeTransformations(it) } + .let { MatchResult(sequenceExpression, it) } .check { checkSmartCastsPreserved(loop, it) } } } @@ -96,10 +112,16 @@ fun match(loop: KtForExpression): ResultTransformationMatch? { return null } - val newState = match.newState + var newState = match.newState // check that old input variable is not needed anymore if (state.inputVariable != newState.inputVariable && state.inputVariable.hasUsages(newState.statements)) return null + if (state.indexVariable != null && match.transformations.any { it.affectsIndex }) { + // index variable is still needed but index in the new sequence is different + if (state.indexVariable!!.hasUsages(newState.statements)) return null + newState = newState.copy(indexVariable = null) + } + sequenceTransformations.addAll(match.transformations) state = newState continue@MatchLoop @@ -111,19 +133,44 @@ fun match(loop: KtForExpression): ResultTransformationMatch? { } //TODO: offer to use of .asSequence() as an option -fun convertLoop(loop: KtForExpression, matchResult: ResultTransformationMatch): KtExpression { - val commentSaver = CommentSaver(matchResult.resultTransformation.commentSavingRange) +fun convertLoop(loop: KtForExpression, matchResult: MatchResult): KtExpression { + val resultTransformation = matchResult.transformationMatch.resultTransformation + val commentSaver = CommentSaver(resultTransformation.commentSavingRange) val callChain = matchResult.generateCallChain(loop) - val result = matchResult.resultTransformation.convertLoop(callChain) + val result = resultTransformation.convertLoop(callChain) - commentSaver.restore(matchResult.resultTransformation.commentRestoringRange(result)) + commentSaver.restore(resultTransformation.commentRestoringRange(result)) return result } -private fun checkSmartCastsPreserved(loop: KtForExpression, matchResult: ResultTransformationMatch): Boolean { +data class LoopData( + val inputVariable: KtCallableDeclaration, + val indexVariable: KtCallableDeclaration?, + val sequenceExpression: KtExpression) + +private fun extractLoopData(loop: KtForExpression): LoopData? { + val loopRange = loop.loopRange ?: return null + + //TODO: recognize other patterns for loop with index + val destructuringParameter = loop.destructuringParameter + if (destructuringParameter != null && destructuringParameter.entries.size == 2) { + val qualifiedExpression = loopRange as? KtDotQualifiedExpression + if (qualifiedExpression != null) { + val call = qualifiedExpression.selectorExpression as? KtCallExpression + //TODO: check that it's the correct "withIndex" + if (call != null && call.calleeExpression.isSimpleName(Name.identifier("withIndex")) && call.valueArguments.isEmpty()) { + return LoopData(destructuringParameter.entries[1], destructuringParameter.entries[0], qualifiedExpression.receiverExpression) + } + } + } + + return LoopData(loop.loopParameter ?: return null, null, loopRange) +} + +private fun checkSmartCastsPreserved(loop: KtForExpression, matchResult: MatchResult): Boolean { val bindingContext = loop.analyze(BodyResolveMode.FULL) val SMARTCAST_KEY = Key("SMARTCAST_KEY") @@ -174,9 +221,9 @@ private fun checkSmartCastsPreserved(loop: KtForExpression, matchResult: ResultT } } -private fun ResultTransformationMatch.generateCallChain(loop: KtForExpression): KtExpression { - var sequenceTransformations = sequenceTransformations - var resultTransformation = resultTransformation +private fun MatchResult.generateCallChain(loop: KtForExpression): KtExpression { + var sequenceTransformations = transformationMatch.sequenceTransformations + var resultTransformation = transformationMatch.resultTransformation while(true) { val last = sequenceTransformations.lastOrNull() ?: break resultTransformation = resultTransformation.mergeWithPrevious(last) ?: break @@ -186,7 +233,7 @@ private fun ResultTransformationMatch.generateCallChain(loop: KtForExpression): val chainCallCount = sequenceTransformations.sumBy { it.chainCallCount } + resultTransformation.chainCallCount val lineBreak = if (chainCallCount > 1) "\n" else "" - var callChain = loop.loopRange!! + var callChain = sequenceExpression val psiFactory = KtPsiFactory(loop) val chainedCallGenerator = object : ChainedCallGenerator { 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 58e568796f7..d3c4a14a2e7 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 @@ -24,6 +24,7 @@ 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.MapIndexedTransformation import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.MapTransformation import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers import org.jetbrains.kotlin.name.FqName @@ -33,14 +34,13 @@ import org.jetbrains.kotlin.renderer.render class AddToCollectionTransformation( loop: KtForExpression, - private val inputVariable: KtCallableDeclaration, private val targetCollection: KtExpression ) : ReplaceLoopResultTransformation(loop) { override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? { return when (previousTransformation) { is FilterTransformation -> { - FilterToTransformation.create(loop, inputVariable, targetCollection, previousTransformation.effectiveCondition()) //TODO: use filterNotTo? + FilterToTransformation.create(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.effectiveCondition()) //TODO: use filterNotTo? } is MapTransformation -> { @@ -80,10 +80,10 @@ class AddToCollectionTransformation( * } */ object Matcher : ResultTransformationMatcher { - override fun match(state: MatchingState): ResultTransformationMatch? { - //TODO: pass indexVariable as null if not used - if (state.indexVariable != null) return null + override val indexVariableUsePossible: Boolean + get() = true + override fun match(state: MatchingState): ResultTransformationMatch? { val statement = state.statements.singleOrNull() ?: return null //TODO: it can be implicit 'this' too val qualifiedExpression = statement as? KtDotQualifiedExpression ?: return null @@ -95,16 +95,22 @@ class AddToCollectionTransformation( val argument = callExpression.valueArguments.singleOrNull() ?: return null val argumentValue = argument.getArgumentExpression() ?: return null - matchWithCollectionInitializationReplacement(state, targetCollection, argumentValue) - ?.let { return it } + if (state.indexVariable == null) { + matchWithCollectionInitializationReplacement(state, targetCollection, argumentValue) + ?.let { return it } + } - val transformation = if (argumentValue.isVariableReference(state.inputVariable)) { - AddToCollectionTransformation(state.outerLoop, state.inputVariable, targetCollection) + if (state.indexVariable == null && argumentValue.isVariableReference(state.inputVariable)) { + return ResultTransformationMatch(AddToCollectionTransformation(state.outerLoop, targetCollection)) + } + else if (state.indexVariable != null) { + val mapIndexedTransformation = MapIndexedTransformation(state.outerLoop, state.inputVariable, state.indexVariable, argumentValue) + val addToCollectionTransformation = AddToCollectionTransformation(state.outerLoop, targetCollection) + return ResultTransformationMatch(addToCollectionTransformation, mapIndexedTransformation) } else { - MapToTransformation.create(state.outerLoop, state.inputVariable, targetCollection, argumentValue) + return ResultTransformationMatch(MapToTransformation.create(state.outerLoop, state.inputVariable, targetCollection, argumentValue)) } - return ResultTransformationMatch(transformation) } private fun matchWithCollectionInitializationReplacement( diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/CountTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/CountTransformation.kt index 7b3167188bf..36a272434c9 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/CountTransformation.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/CountTransformation.kt @@ -71,6 +71,9 @@ class CountTransformation( * } */ object Matcher : ResultTransformationMatcher { + override val indexVariableUsePossible: Boolean + get() = false + override fun match(state: MatchingState): ResultTransformationMatch? { val statement = state.statements.singleOrNull() as? KtUnaryExpression ?: return null if (statement.operationToken != KtTokens.PLUSPLUS) return null 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 ea055bca560..ffb77be4113 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 @@ -71,10 +71,10 @@ class FindAndAssignTransformation( * } */ object Matcher : ResultTransformationMatcher { - override fun match(state: MatchingState): ResultTransformationMatch? { - //TODO: pass indexVariable as null if not used - if (state.indexVariable != null) return null + override val indexVariableUsePossible: Boolean + get() = false + override fun match(state: MatchingState): ResultTransformationMatch? { when (state.statements.size) { 1 -> {} 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 592bd413418..90ec97dcb0c 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 @@ -70,10 +70,10 @@ class FindAndReturnTransformation( * return ... */ object Matcher : ResultTransformationMatcher { - override fun match(state: MatchingState): ResultTransformationMatch? { - //TODO: pass indexVariable as null if not used - if (state.indexVariable != null) return null + override val indexVariableUsePossible: Boolean + get() = false + override fun match(state: MatchingState): ResultTransformationMatch? { val returnInLoop = state.statements.singleOrNull() as? KtReturnExpression ?: return null val returnAfterLoop = state.outerLoop.nextStatement() as? KtReturnExpression ?: return null if (returnInLoop.getLabelName() != returnAfterLoop.getLabelName()) return null 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 1444492f7e4..b5c6ba27e23 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 @@ -69,22 +69,20 @@ class FilterTransformation( */ object Matcher : SequenceTransformationMatcher { override fun match(state: MatchingState): SequenceTransformationMatch? { - if (state.indexVariable != null) return null //TODO? - val ifStatement = state.statements.firstOrNull() as? KtIfExpression ?: return null if (ifStatement.`else` != null) return null val condition = ifStatement.condition ?: return null val then = ifStatement.then ?: return null if (state.statements.size == 1) { - val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, condition, isInverse = false) + val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, state.indexVariable, condition, isInverse = false) val newState = state.copy(statements = listOf(then)) return SequenceTransformationMatch(transformation, newState) } else { val continueExpression = then.blockExpressionsOrSingle().singleOrNull() as? KtContinueExpression ?: return null if (!continueExpression.isBreakOrContinueOfLoop(state.innerLoop)) return null - val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, condition, isInverse = true) + val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, state.indexVariable, condition, isInverse = true) val newState = state.copy(statements = state.statements.drop(1)) return SequenceTransformationMatch(transformation, newState) } @@ -94,11 +92,17 @@ class FilterTransformation( private fun createFilterTransformation( loop: KtForExpression, inputVariable: KtCallableDeclaration, + indexVariable: KtCallableDeclaration?, condition: KtExpression, - isInverse: Boolean): SequenceTransformation { + isInverse: Boolean + ): SequenceTransformation { val effectiveCondition = if (isInverse) condition.negate() else condition + if (indexVariable != null && indexVariable.hasUsages(condition)) { + return FilterIndexedTransformation(loop, inputVariable, indexVariable, effectiveCondition) + } + if (effectiveCondition is KtIsExpression && !effectiveCondition.isNegated && effectiveCondition.leftHandSide.isSimpleName(inputVariable.nameAsSafeName) // we cannot use isVariableReference here because expression can be non-physical @@ -150,3 +154,25 @@ class FilterNotNullTransformation(override val loop: KtForExpression) : Sequence return chainedCallGenerator.generate("filterNotNull()") } } + +class FilterIndexedTransformation( + override val loop: KtForExpression, + val inputVariable: KtCallableDeclaration, + val indexVariable: KtCallableDeclaration, + val condition: KtExpression +) : SequenceTransformation { + + //TODO: how to handle multiple if's using index? + + override val affectsIndex: Boolean + get() = true + + override val presentation: String + get() = "filterIndexed{}" + + override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression { + val lambda = generateLambda(condition, indexVariable, inputVariable) + return chainedCallGenerator.generate("filterIndexed $0:'{}'", lambda) + } +} + 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 82d187ad9c0..3ca2bb073d4 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 @@ -20,9 +20,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.analyze import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade import org.jetbrains.kotlin.idea.intentions.loopToCallChain.* import org.jetbrains.kotlin.idea.util.FuzzyType -import org.jetbrains.kotlin.psi.KtCallableDeclaration -import org.jetbrains.kotlin.psi.KtExpression -import org.jetbrains.kotlin.psi.KtForExpression +import org.jetbrains.kotlin.psi.* import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode class FlatMapTransformation( @@ -53,8 +51,6 @@ class FlatMapTransformation( */ object Matcher : SequenceTransformationMatcher { override fun match(state: MatchingState): SequenceTransformationMatch? { - if (state.indexVariable != null) return null - val nestedLoop = state.statements.singleOrNull() as? KtForExpression ?: return null val transform = nestedLoop.loopRange ?: return null @@ -65,8 +61,21 @@ class FlatMapTransformation( if (iterableType.checkIsSuperTypeOf(nestedSequenceType) == null) return null val nestedLoopBody = nestedLoop.body ?: return null - val newWorkingVariable = nestedLoop.loopParameter ?: return null + + if (state.indexVariable != null && state.indexVariable.hasUsages(transform)) { + // if nested loop range uses index, convert to "mapIndexed {...}.flatMap { it }" + val mapIndexedTransformation = MapIndexedTransformation(state.outerLoop, state.inputVariable, state.indexVariable, transform) + val inputVarExpression = KtPsiFactory(nestedLoop).createExpressionByPattern("$0", state.inputVariable.nameAsSafeName) + val flatMapTransformation = FlatMapTransformation(state.outerLoop, state.inputVariable, inputVarExpression) + val newState = state.copy( + innerLoop = nestedLoop, + statements = listOf(nestedLoopBody), + inputVariable = newWorkingVariable + ) + return SequenceTransformationMatch(listOf(mapIndexedTransformation, flatMapTransformation), newState) + } + val transformation = FlatMapTransformation(state.outerLoop, state.inputVariable, transform) val newState = state.copy( innerLoop = nestedLoop, 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 d5248e5be1f..be63b828d4c 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 @@ -48,16 +48,36 @@ class MapTransformation( */ object Matcher : SequenceTransformationMatcher { override fun match(state: MatchingState): SequenceTransformationMatch? { - if (state.indexVariable != null) return null //TODO? - val declaration = state.statements.firstOrNull() as? KtProperty ?: return null //TODO: support multi-variables val initializer = declaration.initializer ?: return null if (declaration.hasWriteUsages()) return null val restStatements = state.statements.drop(1) - val transformation = MapTransformation(state.outerLoop, state.inputVariable, initializer) + val transformation = if (state.indexVariable != null && state.indexVariable.hasUsages(initializer)) + MapIndexedTransformation(state.outerLoop, state.inputVariable, state.indexVariable, initializer) + else + MapTransformation(state.outerLoop, state.inputVariable, initializer) val newState = state.copy(statements = restStatements, inputVariable = declaration) return SequenceTransformationMatch(transformation, newState) } } -} \ No newline at end of file +} + +class MapIndexedTransformation( + override val loop: KtForExpression, + val inputVariable: KtCallableDeclaration, + val indexVariable: KtCallableDeclaration, + val mapping: KtExpression +) : SequenceTransformation { + + override val affectsIndex: Boolean + get() = false + + override val presentation: String + get() = "mapIndexed{}" + + override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression { + val lambda = generateLambda(mapping, indexVariable, inputVariable) + return chainedCallGenerator.generate("mapIndexed $0:'{}'", lambda) + } +} 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 ad159a11d27..09faddbfb79 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt @@ -81,6 +81,25 @@ fun generateLambda(inputVariable: KtCallableDeclaration, expression: KtExpressio return psiFactory.createExpressionByPattern("{ $0 }", lambdaExpression.bodyExpression!!) as KtLambdaExpression } +fun generateLambda(expression: KtExpression, vararg inputVariables: KtCallableDeclaration): KtLambdaExpression { + return KtPsiFactory(expression).buildExpression { + appendFixedText("{") + + for ((index, variable) in inputVariables.withIndex()) { + if (index > 0) { + appendFixedText(",") + } + appendName(variable.nameAsSafeName) + } + + appendFixedText("->") + + appendExpression(expression) + + appendFixedText("}") + } as KtLambdaExpression +} + fun KtExpression?.isTrueConstant() = this != null && node?.elementType == KtNodeTypes.BOOLEAN_CONSTANT && text == "true" @@ -100,12 +119,14 @@ fun KtCallableDeclaration.hasUsages(inElement: KtElement): Boolean { } fun KtCallableDeclaration.hasUsages(inElements: Collection): Boolean { + assert(this.isPhysical) // TODO: it's a temporary workaround about strange dead-lock when running inspections return inElements.any { ReferencesSearch.search(this, LocalSearchScope(it)).any() } // return ReferencesSearch.search(this, LocalSearchScope(inElements.toTypedArray())).any() } fun KtProperty.hasWriteUsages(): Boolean { + assert(this.isPhysical) if (!isVar) return false return ReferencesSearch.search(this, useScope).any { (it as? KtSimpleNameReference)?.element?.readWriteAccess(useResolveForReadWrite = true)?.isWrite == true diff --git a/idea/testData/intentions/loopToCallChain/any_indexNeeded.kt b/idea/testData/intentions/loopToCallChain/any_indexNeeded.kt new file mode 100644 index 00000000000..97c2d6660f4 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/any_indexNeeded.kt @@ -0,0 +1,11 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'filterIndexed{}.any()'" +fun foo(list: List) { + var found = false + for ((index, s) in list.withIndex()) { + if (s.length > index) { + found = true + break + } + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/any_indexNeeded.kt.after b/idea/testData/intentions/loopToCallChain/any_indexNeeded.kt.after new file mode 100644 index 00000000000..651df60094f --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/any_indexNeeded.kt.after @@ -0,0 +1,7 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'filterIndexed{}.any()'" +fun foo(list: List) { + val found = list + .filterIndexed { index, s -> s.length > index } + .any() +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/filterIndexed.kt b/idea/testData/intentions/loopToCallChain/filterIndexed.kt new file mode 100644 index 00000000000..c9a14c17a1d --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/filterIndexed.kt @@ -0,0 +1,10 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'filterIndexed{}.firstOrNull()'" +fun foo(list: List): String? { + for ((index, s) in list.withIndex()) { + if (s.length > index) { + return s + } + } + return null +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/filterIndexed.kt.after b/idea/testData/intentions/loopToCallChain/filterIndexed.kt.after new file mode 100644 index 00000000000..3b48882aa53 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/filterIndexed.kt.after @@ -0,0 +1,7 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'filterIndexed{}.firstOrNull()'" +fun foo(list: List): String? { + return list + .filterIndexed { index, s -> s.length > index } + .firstOrNull() +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/flatMap_indexUsed.kt b/idea/testData/intentions/loopToCallChain/flatMap_indexUsed.kt new file mode 100644 index 00000000000..a36d6c19b08 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/flatMap_indexUsed.kt @@ -0,0 +1,10 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapIndexed{}.flatMap{}.firstOrNull{}'" +fun foo(list: List): String? { + for ((index, s) in list.withIndex()) { + for (line in s.lines().take(index)) { + if (line.isNotBlank()) return line + } + } + return null +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/flatMap_indexUsed.kt.after b/idea/testData/intentions/loopToCallChain/flatMap_indexUsed.kt.after new file mode 100644 index 00000000000..fe00ed08f1f --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/flatMap_indexUsed.kt.after @@ -0,0 +1,8 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapIndexed{}.flatMap{}.firstOrNull{}'" +fun foo(list: List): String? { + return list + .mapIndexed { index, s -> s.lines().take(index) } + .flatMap { it } + .firstOrNull { it.isNotBlank() } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapAndFilterIndexed.kt b/idea/testData/intentions/loopToCallChain/mapAndFilterIndexed.kt new file mode 100644 index 00000000000..2e98613ef63 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapAndFilterIndexed.kt @@ -0,0 +1,11 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'map{}.filterIndexed{}.firstOrNull()'" +fun foo(list: List): Int? { + for ((index, s) in list.withIndex()) { + val l = s.length + if (l > index) { + return l + } + } + return null +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapAndFilterIndexed.kt.after b/idea/testData/intentions/loopToCallChain/mapAndFilterIndexed.kt.after new file mode 100644 index 00000000000..e8a46c55572 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapAndFilterIndexed.kt.after @@ -0,0 +1,8 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'map{}.filterIndexed{}.firstOrNull()'" +fun foo(list: List): Int? { + return list + .map { it.length } + .filterIndexed { index, l -> l > index } + .firstOrNull() +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapIndexed.kt b/idea/testData/intentions/loopToCallChain/mapIndexed.kt new file mode 100644 index 00000000000..114e772d81c --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapIndexed.kt @@ -0,0 +1,9 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapIndexed{}.firstOrNull{}'" +fun foo(list: List): Int? { + for ((index, s) in list.withIndex()) { + val x = s.length * index + if (x > 0) return x + } + return null +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapIndexed.kt.after b/idea/testData/intentions/loopToCallChain/mapIndexed.kt.after new file mode 100644 index 00000000000..ff38c362f28 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapIndexed.kt.after @@ -0,0 +1,7 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapIndexed{}.firstOrNull{}'" +fun foo(list: List): Int? { + return list + .mapIndexed { index, s -> s.length * index } + .firstOrNull { it > 0 } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapIndexed_twice.kt b/idea/testData/intentions/loopToCallChain/mapIndexed_twice.kt new file mode 100644 index 00000000000..4fe5ca4ac52 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapIndexed_twice.kt @@ -0,0 +1,10 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapIndexed{}.mapIndexed{}.firstOrNull{}'" +fun foo(list: List): Int? { + for ((index, s) in list.withIndex()) { + val x = s.length * index + val y = x + index + if (y > 0) return y + } + return null +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapIndexed_twice.kt.after b/idea/testData/intentions/loopToCallChain/mapIndexed_twice.kt.after new file mode 100644 index 00000000000..f506e4ba2cd --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapIndexed_twice.kt.after @@ -0,0 +1,8 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with 'mapIndexed{}.mapIndexed{}.firstOrNull{}'" +fun foo(list: List): Int? { + return list + .mapIndexed { index, s -> s.length * index } + .mapIndexed { index, x -> x + index } + .firstOrNull { it > 0 } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt b/idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt new file mode 100644 index 00000000000..6bcad90c726 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt @@ -0,0 +1,7 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with '+= mapIndexed{}'" +fun foo(list: List, target: MutableList) { + for ((index, s) in list.withIndex()) { + target.add(s.hashCode() * index) + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt.after b/idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt.after new file mode 100644 index 00000000000..fb139ec672e --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapTo_indexUsed.kt.after @@ -0,0 +1,5 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Replace with '+= mapIndexed{}'" +fun foo(list: List, target: MutableList) { + target += list.mapIndexed { index, s -> s.hashCode() * index } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/mapUsesOldIndexAfterFilter.kt b/idea/testData/intentions/loopToCallChain/mapUsesOldIndexAfterFilter.kt new file mode 100644 index 00000000000..6f993868fd6 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/mapUsesOldIndexAfterFilter.kt @@ -0,0 +1,10 @@ +// WITH_RUNTIME +// IS_APPLICABLE: false +fun foo(list: List): Int? { + for ((index, s) in list.withIndex()) { + if (s.isBlank()) continue + val x = s.length * index + if (x > 1000) return x + } + return null +} \ No newline at end of file