diff --git a/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinHighlightExitPointsHandlerFactory.kt b/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinHighlightExitPointsHandlerFactory.kt index 84e3c9d5750..5af88d5a679 100644 --- a/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinHighlightExitPointsHandlerFactory.kt +++ b/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinHighlightExitPointsHandlerFactory.kt @@ -30,28 +30,51 @@ import org.jetbrains.kotlin.idea.references.mainReference import org.jetbrains.kotlin.lexer.KtTokens import org.jetbrains.kotlin.psi.* import org.jetbrains.kotlin.psi.psiUtil.parents +import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsResultOfLambda import org.jetbrains.kotlin.resolve.inline.InlineUtil class KotlinHighlightExitPointsHandlerFactory : HighlightUsagesHandlerFactoryBase() { companion object { private val RETURN_AND_THROW = TokenSet.create(KtTokens.RETURN_KEYWORD, KtTokens.THROW_KEYWORD) - } - override fun createHighlightUsagesHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? { - if (target is LeafPsiElement && (target.elementType in RETURN_AND_THROW)) { + private fun getOnReturnOrThrowUsageHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? { + if (target !is LeafPsiElement || target.elementType !in RETURN_AND_THROW) { + return null + } + val returnOrThrow = PsiTreeUtil.getParentOfType( target, KtReturnExpression::class.java, KtThrowExpression::class.java ) ?: return null - return MyHandler(editor, file, returnOrThrow) + return OnExitUsagesHandler(editor, file, returnOrThrow) + } + + private fun getOnLambdaCallUsageHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? { + if (target !is LeafPsiElement || target.elementType != KtTokens.IDENTIFIER) { + return null + } + + val refExpr = target.parent as? KtNameReferenceExpression ?: return null + val call = refExpr.parent as? KtCallExpression ?: return null + if (call.calleeExpression != refExpr) return null + + val lambda = call.lambdaArguments.singleOrNull() ?: return null + val literal = lambda.getLambdaExpression()?.functionLiteral ?: return null + + return OnExitUsagesHandler(editor, file, literal) } - return null } - private class MyHandler(editor: Editor, file: PsiFile, val target: KtExpression) : + override fun createHighlightUsagesHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? { + return getOnReturnOrThrowUsageHandler(editor, file, target) + ?: getOnLambdaCallUsageHandler(editor, file, target) + } + + private class OnExitUsagesHandler(editor: Editor, file: PsiFile, val target: KtExpression) : HighlightUsagesHandlerBase(editor, file) { + override fun getTargets() = listOf(target) override fun selectTargets(targets: MutableList, selectionConsumer: Consumer>) { @@ -59,12 +82,48 @@ class KotlinHighlightExitPointsHandlerFactory : HighlightUsagesHandlerFactoryBas } override fun computeUsages(targets: MutableList?) { - val relevantFunction = target.getRelevantDeclaration() + val relevantFunction: KtDeclarationWithBody? = + if (target is KtFunctionLiteral) { + target + } else { + target.getRelevantDeclaration() + } + relevantFunction?.accept(object : KtVisitorVoid() { override fun visitKtElement(element: KtElement) { element.acceptChildren(this) } + override fun visitExpression(expression: KtExpression) { + if (relevantFunction is KtFunctionLiteral) { + if (occurrenceForFunctionLiteralReturnExpression(expression)) { + return + } + } + + super.visitExpression(expression) + } + + private fun occurrenceForFunctionLiteralReturnExpression(expression: KtExpression): Boolean { + if (!KtPsiUtil.isStatement(expression)) return false + + if (expression is KtIfExpression || expression is KtWhenExpression || expression is KtBlockExpression) { + return false + } + + val bindingContext = expression.analyze() + if (!expression.isUsedAsResultOfLambda(bindingContext)) { + return false + } + + if (expression.getRelevantDeclaration() != relevantFunction) { + return false + } + + addOccurrence(expression) + return true + } + private fun visitReturnOrThrow(expression: KtExpression) { if (expression.getRelevantDeclaration() == relevantFunction) { addOccurrence(expression) @@ -85,18 +144,28 @@ class KotlinHighlightExitPointsHandlerFactory : HighlightUsagesHandlerFactoryBas private fun KtExpression.getRelevantDeclaration(): KtDeclarationWithBody? { if (this is KtReturnExpression) { - (this.getTargetLabel()?.mainReference?.resolve() as? KtFunction)?.let { return it } - } - - for (parent in parents) { - if (parent is KtDeclarationWithBody) { - when { - parent is KtPropertyAccessor -> return parent - InlineUtil.canBeInlineArgument(parent) && !InlineUtil.isInlinedArgument(parent as KtFunction, parent.analyze(), false) -> - return parent - } + (this.getTargetLabel()?.mainReference?.resolve() as? KtFunction)?.let { + return it } } - return null + if (this is KtThrowExpression || this is KtReturnExpression) { + for (parent in parents) { + if (parent is KtDeclarationWithBody) { + if (parent is KtPropertyAccessor) { + return parent + } + + if (InlineUtil.canBeInlineArgument(parent) && + !InlineUtil.isInlinedArgument(parent as KtFunction, parent.analyze(), false) + ) { + return parent + } + } + } + + return null + } + + return parents.filterIsInstance().firstOrNull() } \ No newline at end of file diff --git a/idea/testData/usageHighlighter/implicitReturnExpressionsFromExplicitReturnInLambdas.kt b/idea/testData/usageHighlighter/implicitReturnExpressionsFromExplicitReturnInLambdas.kt new file mode 100644 index 00000000000..33e341b7d66 --- /dev/null +++ b/idea/testData/usageHighlighter/implicitReturnExpressionsFromExplicitReturnInLambdas.kt @@ -0,0 +1,17 @@ +fun some(a: Int, b: Int) { + run { + val i = 12 + val j = 13 + if (a > 50) { + if (b > 100) { + i + j + } else { + i * j + } + } else { + ~return@run false + } + } +} + +fun run(a: () -> T) {} \ No newline at end of file diff --git a/idea/testData/usageHighlighter/implicitReturnExpressionsInLambdasNoHightlighting.kt b/idea/testData/usageHighlighter/implicitReturnExpressionsInLambdasNoHightlighting.kt new file mode 100644 index 00000000000..9e541e39396 --- /dev/null +++ b/idea/testData/usageHighlighter/implicitReturnExpressionsInLambdasNoHightlighting.kt @@ -0,0 +1,17 @@ +fun some(a: Int, b: Int) { + run { + val i = 12 + val j = 13 + if (a > 50) { + if (b > 100) { + 101 + j + } else { + ~102 * j + } + } else { + return@run false + } + } +} + +fun run(a: () -> T) {} \ No newline at end of file diff --git a/idea/testData/usageHighlighter/lambdaCallReturnExpressions.kt b/idea/testData/usageHighlighter/lambdaCallReturnExpressions.kt new file mode 100644 index 00000000000..171ce149528 --- /dev/null +++ b/idea/testData/usageHighlighter/lambdaCallReturnExpressions.kt @@ -0,0 +1,17 @@ +fun some(a: Int, b: Int) { + ~run { + val i = 12 + val j = 13 + if (a > 50) { + if (b > 100) { + i + j + } else { + i * j + } + } else { + return@run false + } + } +} + +fun run(a: () -> T) {} diff --git a/idea/testData/usageHighlighter/lambdaCallReturnExpressionsInline.kt b/idea/testData/usageHighlighter/lambdaCallReturnExpressionsInline.kt new file mode 100644 index 00000000000..a6d7d897e12 --- /dev/null +++ b/idea/testData/usageHighlighter/lambdaCallReturnExpressionsInline.kt @@ -0,0 +1,18 @@ +fun some(a: Int, b: Int) { + ~run { + val i = 12 + val j = 13 + if (a > 50) { + if (b > 100) { + i + j + } else { + return@some + } + } else { + return@run false + } + } +} + +inline fun run(a: () -> T) { +} \ No newline at end of file diff --git a/idea/testData/usageHighlighter/lambdaCallReturnExpressionsInlineUnit.kt b/idea/testData/usageHighlighter/lambdaCallReturnExpressionsInlineUnit.kt new file mode 100644 index 00000000000..47fddc447df --- /dev/null +++ b/idea/testData/usageHighlighter/lambdaCallReturnExpressionsInlineUnit.kt @@ -0,0 +1,18 @@ +fun some(a: Int, b: Int) { + ~run { + val i = 12 + val j = 13 + if (a > 50) { + if (b > 100) { + i + j + } else { + return@some + } + } else { + return@run + } + } +} + +inline fun run(a: () -> Unit) { +} \ No newline at end of file diff --git a/idea/tests/org/jetbrains/kotlin/idea/highlighter/UsageHighlightingTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/highlighter/UsageHighlightingTestGenerated.java index 7ba5c2a10c0..89b3fe95f1b 100644 --- a/idea/tests/org/jetbrains/kotlin/idea/highlighter/UsageHighlightingTestGenerated.java +++ b/idea/tests/org/jetbrains/kotlin/idea/highlighter/UsageHighlightingTestGenerated.java @@ -34,6 +34,16 @@ public class UsageHighlightingTestGenerated extends AbstractUsageHighlightingTes runTest("idea/testData/usageHighlighter/implicitIt.kt"); } + @TestMetadata("implicitReturnExpressionsFromExplicitReturnInLambdas.kt") + public void testImplicitReturnExpressionsFromExplicitReturnInLambdas() throws Exception { + runTest("idea/testData/usageHighlighter/implicitReturnExpressionsFromExplicitReturnInLambdas.kt"); + } + + @TestMetadata("implicitReturnExpressionsInLambdasNoHightlighting.kt") + public void testImplicitReturnExpressionsInLambdasNoHightlighting() throws Exception { + runTest("idea/testData/usageHighlighter/implicitReturnExpressionsInLambdasNoHightlighting.kt"); + } + @TestMetadata("importAlias.kt") public void testImportAlias() throws Exception { runTest("idea/testData/usageHighlighter/importAlias.kt"); @@ -54,6 +64,21 @@ public class UsageHighlightingTestGenerated extends AbstractUsageHighlightingTes runTest("idea/testData/usageHighlighter/labeledLoop.kt"); } + @TestMetadata("lambdaCallReturnExpressions.kt") + public void testLambdaCallReturnExpressions() throws Exception { + runTest("idea/testData/usageHighlighter/lambdaCallReturnExpressions.kt"); + } + + @TestMetadata("lambdaCallReturnExpressionsInline.kt") + public void testLambdaCallReturnExpressionsInline() throws Exception { + runTest("idea/testData/usageHighlighter/lambdaCallReturnExpressionsInline.kt"); + } + + @TestMetadata("lambdaCallReturnExpressionsInlineUnit.kt") + public void testLambdaCallReturnExpressionsInlineUnit() throws Exception { + runTest("idea/testData/usageHighlighter/lambdaCallReturnExpressionsInlineUnit.kt"); + } + @TestMetadata("localVal.kt") public void testLocalVal() throws Exception { runTest("idea/testData/usageHighlighter/localVal.kt");