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 2e877427d17..25cbb5c7aeb 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 @@ -27,7 +27,7 @@ import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange class FindAndAssignTransformation( private val loop: KtForExpression, override val inputVariable: KtCallableDeclaration, - private val stdlibFunName: String, + private val generator: (chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?) -> KtExpression, private val initialDeclaration: KtProperty, private val filter: KtExpression? = null ) : ResultTransformation { @@ -35,20 +35,14 @@ class FindAndAssignTransformation( override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? { if (previousTransformation !is FilterTransformation) return null assert(filter == null) { "Should not happen because no 2 consecutive FilterTransformation's possible"} - return FindAndAssignTransformation(loop, previousTransformation.inputVariable, stdlibFunName, initialDeclaration, previousTransformation.buildRealCondition()) + return FindAndAssignTransformation(loop, previousTransformation.inputVariable, generator, initialDeclaration, previousTransformation.buildRealCondition()) } override val commentSavingRange = PsiChildRange(initialDeclaration, loop.unwrapIfLabeled()) override val commentRestoringRange = commentSavingRange.withoutLastStatement() override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression { - return if (filter == null) { - chainedCallGenerator.generate("$stdlibFunName()") - } - else { - val lambda = generateLambda(inputVariable, filter) - chainedCallGenerator.generate("$stdlibFunName $0:'{}'", lambda) - } + return generator(chainedCallGenerator, filter) } override fun convertLoop(resultCallChain: KtExpression): KtExpression { @@ -107,9 +101,9 @@ class FindAndAssignTransformation( val usageCountInLoop = ReferencesSearch.search(declarationBeforeLoop, LocalSearchScope(state.outerLoop)).count() if (usageCountInLoop != 1) return null // this should be the only usage of this variable inside the loop - val stdlibFunName = stdlibFunNameForFind(right, initializer, state.workingVariable, findFirst) ?: return null + val generator = buildFindOperationGenerator(right, initializer, state.workingVariable, findFirst) ?: return null - val transformation = FindAndAssignTransformation(state.outerLoop, state.workingVariable, stdlibFunName, declarationBeforeLoop) + val transformation = FindAndAssignTransformation(state.outerLoop, state.workingVariable, generator, declarationBeforeLoop) return ResultTransformationMatch(transformation) } } 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 d88008d2e80..c1b636c7d0b 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 @@ -27,7 +27,7 @@ import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange class FindAndReturnTransformation( private val loop: KtForExpression, override val inputVariable: KtCallableDeclaration, - private val stdlibFunName: String, + private val generator: (chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?) -> KtExpression, private val endReturn: KtReturnExpression, private val filter: KtExpression? = null ) : ResultTransformation { @@ -35,7 +35,7 @@ class FindAndReturnTransformation( override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? { if (previousTransformation !is FilterTransformation) return null assert(filter == null) { "Should not happen because no 2 consecutive FilterTransformation's possible"} - return FindAndReturnTransformation(loop, previousTransformation.inputVariable, stdlibFunName, endReturn, previousTransformation.buildRealCondition()) + return FindAndReturnTransformation(loop, previousTransformation.inputVariable, generator, endReturn, previousTransformation.buildRealCondition()) } override val commentSavingRange = PsiChildRange(loop.unwrapIfLabeled(), endReturn) @@ -43,13 +43,7 @@ class FindAndReturnTransformation( override val commentRestoringRange = commentSavingRange.withoutFirstStatement() override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression { - return if (filter == null) { - chainedCallGenerator.generate("$stdlibFunName()") - } - else { - val lambda = generateLambda(inputVariable, filter) - chainedCallGenerator.generate("$stdlibFunName $0:'{}'", lambda) - } + return generator(chainedCallGenerator, filter) } override fun convertLoop(resultCallChain: KtExpression): KtExpression { @@ -78,9 +72,9 @@ class FindAndReturnTransformation( val returnValueInLoop = returnInLoop.returnedExpression ?: return null val returnValueAfterLoop = returnAfterLoop.returnedExpression ?: return null - val stdlibFunName = stdlibFunNameForFind(returnValueInLoop, returnValueAfterLoop, state.workingVariable, findFirst = true) ?: return null + val generator = buildFindOperationGenerator(returnValueInLoop, returnValueAfterLoop, state.workingVariable, findFirst = true) ?: return null - val transformation = FindAndReturnTransformation(state.outerLoop, state.workingVariable, stdlibFunName, returnAfterLoop) + val transformation = FindAndReturnTransformation(state.outerLoop, state.workingVariable, generator, returnAfterLoop) return ResultTransformationMatch(transformation) } } 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 cc502fe90c4..f5e3266ce24 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt @@ -34,6 +34,7 @@ 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.constants.evaluate.ConstantExpressionEvaluator import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull @@ -95,16 +96,48 @@ fun KtProperty.hasWriteUsages(): Boolean { } } -fun stdlibFunNameForFind(valueIfFound: KtExpression, valueIfNotFound: KtExpression, workingVariable: KtCallableDeclaration, findFirst: Boolean): String? { - return when { - valueIfNotFound.isNullExpression() && valueIfFound.isVariableReference(workingVariable) -> if (findFirst) "firstOrNull" else "lastOrNull" +fun buildFindOperationGenerator( + valueIfFound: KtExpression, + valueIfNotFound: KtExpression, + workingVariable: KtCallableDeclaration, + findFirst: Boolean +): ((chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?) -> KtExpression)? { + assert(valueIfFound.isPhysical) + assert(valueIfNotFound.isPhysical) + + fun generateChainedCall(stdlibFunName: String, chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?): KtExpression { + return if (filter == null) { + chainedCallGenerator.generate("$stdlibFunName()") + } + else { + val lambda = generateLambda(workingVariable, filter) + chainedCallGenerator.generate("$stdlibFunName $0:'{}'", lambda) + } + } + + + val stdlibFunName = when { + valueIfNotFound.isNullExpression() && valueIfFound.isVariableReference(workingVariable) -> if (findFirst) "firstOrNull" else "lastOrNull" //TODO: ?: if not null valueIfFound.isTrueConstant() && valueIfNotFound.isFalseConstant() -> "any" valueIfFound.isFalseConstant() && valueIfNotFound.isTrueConstant() -> "none" - else -> /*TODO: allow other constants*/ null + workingVariable.hasUsages(listOf(valueIfFound)) -> /*TODO*/ return null + + // initial value is compile-time constant + ConstantExpressionEvaluator.getConstant(valueIfNotFound, valueIfNotFound.analyze(BodyResolveMode.PARTIAL)) != null -> { + return { chainedCallGenerator, filter -> + val chainedCall = generateChainedCall("any", chainedCallGenerator, filter) + KtPsiFactory(chainedCall).createExpressionByPattern("if ($0) $1 else $2", chainedCall, valueIfFound, valueIfNotFound) + } + + } + + else -> return null } + + return { chainedCallGenerator, filter -> generateChainedCall(stdlibFunName, chainedCallGenerator, filter) } } fun KtExpressionWithLabel.isBreakOrContinueOfLoop(loop: KtLoopExpression): Boolean { diff --git a/idea/testData/intentions/loopToCallChain/any_ifAssign_nonBooleanResult.kt b/idea/testData/intentions/loopToCallChain/any_ifAssign_nonBooleanResult.kt new file mode 100644 index 00000000000..17e7e736b41 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/any_ifAssign_nonBooleanResult.kt @@ -0,0 +1,10 @@ +// WITH_RUNTIME +fun foo(list: List) { + var result = 0 + for (s in list) { + if (s.length > 0) { + result = 1 + break + } + } +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/any_ifAssign_nonBooleanResult.kt.after b/idea/testData/intentions/loopToCallChain/any_ifAssign_nonBooleanResult.kt.after new file mode 100644 index 00000000000..bf40f90df27 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/any_ifAssign_nonBooleanResult.kt.after @@ -0,0 +1,4 @@ +// WITH_RUNTIME +fun foo(list: List) { + val result = if (list.any { it.length > 0 }) 1 else 0 +} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/any_ifAssign_nonBooleanResult2.kt b/idea/testData/intentions/loopToCallChain/any_ifAssign_nonBooleanResult2.kt new file mode 100644 index 00000000000..a72376cae9a --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/any_ifAssign_nonBooleanResult2.kt @@ -0,0 +1,13 @@ +// WITH_RUNTIME +// IS_APPLICABLE: false +fun foo(list: List) { + var result = takeInt() + for (s in list) { + if (s.length > 0) { + result = 1 + break + } + } +} + +fun takeInt(): Int = 0 \ No newline at end of file