Show lambda exist points when cursor is at the correspondent call (KT-29381)

#KT-29381 Fixed
2018-12-28 19:28:47 +03:00
parent 5ec7643b5a
commit a0b8980d77
7 changed files with 199 additions and 18 deletions
@@ -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<KtExpression>(
                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<PsiElement>(editor, file) {
+
        override fun getTargets() = listOf(target)

        override fun selectTargets(targets: MutableList<PsiElement>, selectionConsumer: Consumer<MutableList<PsiElement>>) {
@@ -59,12 +82,48 @@ class KotlinHighlightExitPointsHandlerFactory : HighlightUsagesHandlerFactoryBas
        }

        override fun computeUsages(targets: MutableList<PsiElement>?) {
-            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<KtDeclarationWithBody>().firstOrNull()
 }
@@ -0,0 +1,17 @@
+fun some(a: Int, b: Int) {
+    run {
+        val i = 12
+        val j = 13
+        if (a > 50) {
+            if (b > 100) {
+                <info descr="null">i + j</info>
+            } else {
+                <info descr="null">i * j</info>
+            }
+        } else {
+            <info descr="null">~return@run false</info>
+        }
+    }
+}
+
+fun <T> run(a: () -> T) {}
@@ -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 {
+                <info descr="null">~102</info> * j
+            }
+        } else {
+            return@run false
+        }
+    }
+}
+
+fun <T> run(a: () -> T) {}
@@ -0,0 +1,17 @@
+fun some(a: Int, b: Int) {
+    ~run {
+        val i = 12
+        val j = 13
+        if (a > 50) {
+            if (b > 100) {
+                <info descr="null">i + j</info>
+            } else {
+                <info descr="null">i * j</info>
+            }
+        } else {
+            <info descr="null">return@run false</info>
+        }
+    }
+}
+
+fun <T> run(a: () -> T) {}
@@ -0,0 +1,18 @@
+fun some(a: Int, b: Int) {
+    ~run {
+        val i = 12
+        val j = 13
+        if (a > 50) {
+            if (b > 100) {
+                <info descr="null">i + j</info>
+            } else {
+                return@some
+            }
+        } else {
+            <info descr="null">return@run false</info>
+        }
+    }
+}
+
+inline fun <T> run(a: () -> T) {
+}
@@ -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 {
+            <info descr="null">return@run</info>
+        }
+    }
+}
+
+inline fun <T> run(a: () -> Unit) {
+}
@@ -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");