diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/LiftReturnOrAssignmentInspection.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/LiftReturnOrAssignmentInspection.kt index 48cae872f83..19f6ee6473c 100644 --- a/idea/src/org/jetbrains/kotlin/idea/inspections/LiftReturnOrAssignmentInspection.kt +++ b/idea/src/org/jetbrains/kotlin/idea/inspections/LiftReturnOrAssignmentInspection.kt @@ -22,10 +22,8 @@ import com.intellij.psi.PsiElement import org.jetbrains.kotlin.idea.intentions.branchedTransformations.BranchedFoldingUtils import org.jetbrains.kotlin.idea.intentions.branchedTransformations.isIfBranch import org.jetbrains.kotlin.idea.intentions.branchedTransformations.lineCount -import org.jetbrains.kotlin.psi.KtExpression -import org.jetbrains.kotlin.psi.KtIfExpression -import org.jetbrains.kotlin.psi.KtVisitorVoid -import org.jetbrains.kotlin.psi.KtWhenExpression +import org.jetbrains.kotlin.psi.* +import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType import org.jetbrains.kotlin.psi.psiUtil.getParentOfType class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() { @@ -36,12 +34,13 @@ class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() { if (expression.lineCount() > LINES_LIMIT) return if (expression.isIfBranch()) return - val returnNumber = BranchedFoldingUtils.getFoldableReturnNumber(expression) - if (returnNumber > 0) { + val foldableReturns = BranchedFoldingUtils.getFoldableReturns(expression) + if (foldableReturns?.isNotEmpty() == true) { + val hasOtherReturns = expression.anyDescendantOfType { it !in foldableReturns } holder.registerProblem( keyword, "Return can be lifted out of '${keyword.text}'", - if (returnNumber > 1) ProblemHighlightType.GENERIC_ERROR_OR_WARNING + if (!hasOtherReturns && foldableReturns.size > 1) ProblemHighlightType.GENERIC_ERROR_OR_WARNING else ProblemHighlightType.INFORMATION, LiftReturnOutFix(keyword.text) ) diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedFoldingUtils.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedFoldingUtils.kt index 0d49047a6a7..9aa0b3be5c4 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedFoldingUtils.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedFoldingUtils.kt @@ -95,44 +95,47 @@ object BranchedFoldingUtils { return assignments.size } - private fun getFoldableReturnNumber(branches: List) = - branches.fold(0) { prevNumber, branch -> - when { - prevNumber == -1 -> -1 - getFoldableBranchedReturn(branch) != null -> prevNumber + 1 - else -> { - val currNumber = getFoldableReturnNumber(branch?.lastBlockStatementOrThis()) - if (currNumber == -1) -1 else prevNumber + currNumber - } + private fun getFoldableReturns(branches: List): List? = + branches.fold?>(mutableListOf()) { prevList, branch -> + if (prevList == null) return@fold null + val foldableBranchedReturn = getFoldableBranchedReturn(branch) + if (foldableBranchedReturn != null) { + prevList.add(foldableBranchedReturn) } + else { + val currReturns = getFoldableReturns(branch?.lastBlockStatementOrThis()) ?: return@fold null + prevList += currReturns + } + prevList } - internal fun getFoldableReturnNumber(expression: KtExpression?): Int = when (expression) { + internal fun getFoldableReturns(expression: KtExpression?): List? = when (expression) { is KtWhenExpression -> { val entries = expression.entries when { - !KtPsiUtil.checkWhenExpressionHasSingleElse(expression) -> -1 - entries.isEmpty() -> -1 - else -> getFoldableReturnNumber(entries.map { it.expression }) + !KtPsiUtil.checkWhenExpressionHasSingleElse(expression) -> null + entries.isEmpty() -> null + else -> getFoldableReturns(entries.map { it.expression }) } } is KtIfExpression -> { val branches = expression.branches when { - branches.isEmpty() -> -1 - branches.lastOrNull()?.getStrictParentOfType()?.`else` == null -> -1 - else -> getFoldableReturnNumber(branches) + branches.isEmpty() -> null + branches.lastOrNull()?.getStrictParentOfType()?.`else` == null -> null + else -> getFoldableReturns(branches) } } is KtCallExpression -> { - if (expression.analyze().getType(expression)?.isNothing() == true) 0 else -1 + if (expression.analyze().getType(expression)?.isNothing() == true) emptyList() else null } - is KtBreakExpression, is KtContinueExpression, is KtThrowExpression -> 0 - else -> -1 + is KtBreakExpression, is KtContinueExpression, is KtThrowExpression -> emptyList() + else -> null } - fun canFoldToReturn(expression: KtExpression?): Boolean = - getFoldableReturnNumber(expression) > 0 + private fun getFoldableReturnNumber(expression: KtExpression?) = getFoldableReturns(expression)?.size ?: -1 + + fun canFoldToReturn(expression: KtExpression?): Boolean = getFoldableReturnNumber(expression) > 0 fun foldToAssignment(expression: KtExpression) { var lhs: KtExpression? = null diff --git a/idea/testData/inspectionsLocal/liftOut/whenToReturn/otherReturns.kt b/idea/testData/inspectionsLocal/liftOut/whenToReturn/otherReturns.kt new file mode 100644 index 00000000000..97a96b4c413 --- /dev/null +++ b/idea/testData/inspectionsLocal/liftOut/whenToReturn/otherReturns.kt @@ -0,0 +1,13 @@ +// HIGHLIGHT: INFORMATION + +fun test(n: Int, arg: String?): String { + when (n) { + 1 -> { + if (arg == null) return "" + return "** $arg" + } + else -> { + return "Strange" + } + } +} \ No newline at end of file diff --git a/idea/testData/inspectionsLocal/liftOut/whenToReturn/otherReturns.kt.after b/idea/testData/inspectionsLocal/liftOut/whenToReturn/otherReturns.kt.after new file mode 100644 index 00000000000..0398d2c4e24 --- /dev/null +++ b/idea/testData/inspectionsLocal/liftOut/whenToReturn/otherReturns.kt.after @@ -0,0 +1,13 @@ +// HIGHLIGHT: INFORMATION + +fun test(n: Int, arg: String?): String { + return when (n) { + 1 -> { + if (arg == null) return "" + "** $arg" + } + else -> { + "Strange" + } + } +} \ No newline at end of file diff --git a/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java index 8f1c834bfed..df1a0ac30a5 100644 --- a/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java +++ b/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java @@ -643,6 +643,12 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest { doTest(fileName); } + @TestMetadata("otherReturns.kt") + public void testOtherReturns() throws Exception { + String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn/otherReturns.kt"); + doTest(fileName); + } + @TestMetadata("simpleWhen.kt") public void testSimpleWhen() throws Exception { String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn/simpleWhen.kt");