KT-5875 Support code completion after "return@"

#KT-5875 Fixed

idea/src/org/jetbrains/jet/plugin/completion/CompletionSession.kt

@@ -70,7 +70,7 @@ abstract class CompletionSessionBase(protected val configuration: CompletionSess
         val reference = position.getParent()?.getReferences()?.firstIsInstanceOrNull<JetSimpleNameReference>()
         if (reference != null) {
             if (reference.expression is JetLabelReferenceExpression) {
-                this.expression = reference.expression.getParent().getParent() as? JetThisExpression
+                this.expression = reference.expression.getParent().getParent() as? JetExpressionWithLabel
                 this.reference = null
             }
             else {
@@ -242,6 +242,13 @@ class BasicCompletionSession(configuration: CompletionSessionConfiguration,
                         }
                     }
 
+                    // if "return" is parsed correctly in the current context - insert it and all return@xxx items
+                    "return" -> {
+                        if (expression != null) {
+                            collector.addElements(returnExpressionItems(bindingContext!!, expression))
+                        }
+                    }
+
                     else -> collector.addElement(lookupElement)
                 }
             }

idea/src/org/jetbrains/jet/plugin/completion/CompletionUtils.kt

@@ -58,6 +58,11 @@ import com.intellij.codeInsight.lookup.LookupElementBuilder
 import org.jetbrains.jet.renderer.DescriptorRenderer
 import org.jetbrains.jet.plugin.util.FuzzyType
 import org.jetbrains.jet.lang.psi.JetLabeledExpression
+import org.jetbrains.jet.lang.psi.JetElement
+import org.jetbrains.jet.lang.psi.psiUtil.parents
+import org.jetbrains.jet.lang.psi.JetReferenceExpression
+import org.jetbrains.jet.lang.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.jet.lang.psi.JetDeclarationWithBody
 
 enum class ItemPriority {
     MULTIPLE_ARGUMENTS_ITEM
@@ -210,22 +215,17 @@ fun shouldCompleteThisItems(prefixMatcher: PrefixMatcher): Boolean {
 
 data class ThisItemInfo(val factory: () -> LookupElement, val type: FuzzyType)
 
-fun thisExpressionItems(bindingContext: BindingContext, context: JetExpression): Collection<ThisItemInfo> {
+fun thisExpressionItems(bindingContext: BindingContext, position: JetExpression): Collection<ThisItemInfo> {
-    val scope = bindingContext[BindingContext.RESOLUTION_SCOPE, context] ?: return listOf()
+    val scope = bindingContext[BindingContext.RESOLUTION_SCOPE, position] ?: return listOf()
 
     val result = ArrayList<ThisItemInfo>()
     for ((i, receiver) in scope.getImplicitReceiversWithInstance().withIndex()) {
         val thisType = receiver.getType()
         val fuzzyType = FuzzyType(thisType, listOf())
-        val qualifier = if (i == 0) null else (thisQualifierName(receiver) ?: continue)
+        val label = if (i == 0) null else (thisQualifierName(receiver) ?: continue)
 
         fun createLookupElement(): LookupElement {
-            var element = LookupElementBuilder.create(KeywordLookupObject, if (qualifier == null) "this" else "this@" + qualifier)
+            var element = createKeywordWithLabelElement("this", label)
-            element = element.withPresentableText("this")
-            element = element.withBoldness(true)
-            if (qualifier != null) {
-                element = element.withTailText("@" + qualifier, false)
-            }
             element = element.withTypeText(DescriptorRenderer.SHORT_NAMES_IN_TYPES.renderType(thisType))
             return element
         }
@@ -243,17 +243,71 @@ private fun thisQualifierName(receiver: ReceiverParameterDescriptor): String? {
     }
 
     val functionLiteral = DescriptorToSourceUtils.descriptorToDeclaration(descriptor) as? JetFunctionLiteral ?: return null
-    val literalParent = (functionLiteral.getParent() as JetFunctionLiteralExpression).getParent()
+    return functionLiteralLabel(functionLiteral)
-    when (literalParent) {
+}
-        is JetLabeledExpression -> return literalParent.getLabelName()
 
-        is JetValueArgument -> {
+private fun functionLiteralLabel(functionLiteral: JetFunctionLiteral): String?
-            val parent = literalParent.getParent()
+        = functionLiteralLabelAndCall(functionLiteral).first
-            val callExpression = (if (parent is JetValueArgumentList) parent else literalParent)
+
-                    .getParent() as? JetCallExpression
+private fun functionLiteralLabelAndCall(functionLiteral: JetFunctionLiteral): Pair<String?, JetCallExpression?> {
-            return (callExpression?.getCalleeExpression() as? JetSimpleNameExpression)?.getReferencedName()
+    val literalParent = (functionLiteral.getParent() as JetFunctionLiteralExpression).getParent()
+
+    fun JetValueArgument.callExpression(): JetCallExpression? {
+        val parent = getParent()
+        return (if (parent is JetValueArgumentList) parent else this).getParent() as? JetCallExpression
+    }
+
+    when (literalParent) {
+        is JetLabeledExpression -> {
+            val callExpression = (literalParent.getParent() as? JetValueArgument)?.callExpression()
+            return Pair(literalParent.getLabelName(), callExpression)
         }
 
-        else -> return null
+        is JetValueArgument -> {
+            val callExpression = literalParent.callExpression()
+            val label = (callExpression?.getCalleeExpression() as? JetSimpleNameExpression)?.getReferencedName()
+            return Pair(label, callExpression)
+        }
+
+        else -> {
+            return Pair(null, null)
+        }
     }
 }
+
+fun returnExpressionItems(bindingContext: BindingContext, position: JetElement): Collection<LookupElement> {
+    val result = ArrayList<LookupElement>()
+    for (parent in position.parents()) {
+        when (parent) {
+            is JetFunctionLiteral -> {
+                val (label, call) = functionLiteralLabelAndCall(parent)
+                if (label != null) {
+                    result.add(createKeywordWithLabelElement("return", label))
+                }
+
+                // check if the current function literal is inlined and stop processing outer declarations if it's not
+                val callee = call?.getCalleeExpression() as? JetReferenceExpression ?: break // not inlined
+                val target = bindingContext[BindingContext.REFERENCE_TARGET, callee] as? SimpleFunctionDescriptor ?: break // not inlined
+                if (!target.getInlineStrategy().isInline()) break // not inlined
+            }
+
+            is JetDeclarationWithBody -> {
+                if (parent.hasBlockBody()) {
+                    result.add(createKeywordWithLabelElement("return", null))
+                }
+                break
+            }
+        }
+    }
+    return result
+}
+
+private fun createKeywordWithLabelElement(keyword: String, label: String?): LookupElementBuilder {
+    var element = LookupElementBuilder.create(KeywordLookupObject, if (label == null) keyword else "$keyword@$label")
+    element = element.withPresentableText(keyword)
+    element = element.withBoldness(true)
+    if (label != null) {
+        element = element.withTailText("@$label", false)
+    }
+    return element
+}

idea/testData/completion/keywords/InFunctionExpressionBody.kt

@@ -9,11 +9,10 @@ fun foo() = <caret>
 // EXIST: null
 // EXIST: object
 // EXIST: package
-// EXIST: return
 // EXIST: super
 // EXIST: throw
 // EXIST: true
 // EXIST: try
 // EXIST: when
 // EXIST: while
-// NUMBER: 16
+// NUMBER: 15

idea/testData/completion/keywords/InGetterExpressionBody.kt

@@ -10,11 +10,10 @@ val prop: Int
 // EXIST: null
 // EXIST: object
 // EXIST: package
-// EXIST: return
 // EXIST: super
 // EXIST: throw
 // EXIST: true
 // EXIST: try
 // EXIST: when
 // EXIST: while
-// NUMBER: 16
+// NUMBER: 15

idea/testData/completion/keywords/InPropertyInitializer.kt

@@ -9,11 +9,10 @@ var a : Int = <caret>
 // EXIST: null
 // EXIST: object
 // EXIST: package
-// EXIST: return
 // EXIST: super
 // EXIST: throw
 // EXIST: true
 // EXIST: try
 // EXIST: when
 // EXIST: while
-// NUMBER: 16
+// NUMBER: 15

idea/testData/completion/keywords/Return1.kt

@@ -0,0 +1,14 @@
+fun foo() {
+    takeHandler1 {
+        takeHandler2({ ret<caret> })
+    }
+}
+
+inline fun takeHandler1(handler: () -> Unit){}
+inline fun takeHandler2(handler: () -> Unit){}
+
+// INVOCATION_COUNT: 1
+// EXIST: { lookupString: "return", itemText: "return", tailText: null, attributes: "bold" }
+// EXIST: { lookupString: "return@takeHandler1", itemText: "return", tailText: "@takeHandler1", attributes: "bold" }
+// EXIST: { lookupString: "return@takeHandler2", itemText: "return", tailText: "@takeHandler2", attributes: "bold" }
+// NUMBER: 3

idea/testData/completion/keywords/Return2.kt

@@ -0,0 +1,13 @@
+fun foo() {
+    takeHandler1 {
+        takeHandler2 { <caret> }
+    }
+}
+
+fun takeHandler1(handler: () -> Unit){}
+inline fun takeHandler2(handler: () -> Unit){}
+
+// INVOCATION_COUNT: 1
+// ABSENT: return
+// EXIST: { lookupString: "return@takeHandler1", itemText: "return", tailText: "@takeHandler1", attributes: "bold" }
+// EXIST: { lookupString: "return@takeHandler2", itemText: "return", tailText: "@takeHandler2", attributes: "bold" }

idea/testData/completion/keywords/Return3.kt

@@ -0,0 +1,13 @@
+fun foo() {
+    takeHandler1 {
+        takeHandler2 { <caret> }
+    }
+}
+
+fun takeHandler1(handler: () -> Unit){}
+fun takeHandler2(handler: () -> Unit){}
+
+// INVOCATION_COUNT: 1
+// ABSENT: return
+// ABSENT: return@takeHandler1
+// EXIST: { lookupString: "return@takeHandler2", itemText: "return", tailText: "@takeHandler2", attributes: "bold" }

idea/testData/completion/keywords/Return4.kt

@@ -0,0 +1,12 @@
+fun foo() {
+    takeHandler @label {
+        <caret>
+    }
+}
+
+fun takeHandler(handler: () -> Unit){}
+
+// INVOCATION_COUNT: 1
+// ABSENT: return
+// ABSENT: "return@takeHandler"
+// EXIST: { lookupString: "return@label", itemText: "return", tailText: "@label", attributes: "bold" }

idea/testData/completion/keywords/Return5.kt

@@ -0,0 +1,12 @@
+fun foo() {
+    takeHandler @label {
+        <caret>
+    }
+}
+
+inline fun takeHandler(handler: () -> Unit){}
+
+// INVOCATION_COUNT: 1
+// EXIST: { lookupString: "return", itemText: "return", tailText: null, attributes: "bold" }
+// ABSENT: "return@takeHandler"
+// EXIST: { lookupString: "return@label", itemText: "return", tailText: "@label", attributes: "bold" }

idea/testData/completion/keywords/Return6.kt

@@ -0,0 +1,13 @@
+fun foo() {
+    takeHandler1 {
+        takeHandler2({ return@<caret> })
+    }
+}
+
+inline fun takeHandler1(handler: () -> Unit){}
+inline fun takeHandler2(handler: () -> Unit){}
+
+// INVOCATION_COUNT: 1
+// EXIST: { lookupString: "return@takeHandler1", itemText: "return", tailText: "@takeHandler1", attributes: "bold" }
+// EXIST: { lookupString: "return@takeHandler2", itemText: "return", tailText: "@takeHandler2", attributes: "bold" }
+// NUMBER: 2

idea/testData/completion/keywords/Return7.kt

@@ -0,0 +1,6 @@
+val property: Int
+    get() {
+        <caret>
+    }
+
+// EXIST: return

idea/testData/completion/keywords/Return8.kt

@@ -0,0 +1,3 @@
+fun foo(p: Int?): Int = p ?: <caret>
+
+// ABSENT: return

idea/testData/completion/keywords/Return9.kt

@@ -0,0 +1,9 @@
+inline fun run (p: () -> Unit) {}
+
+fun foo() = run {
+    <caret>
+}
+
+// INVOCATION_COUNT: 1
+// EXIST: { lookupString: "return@run", itemText: "return", tailText: "@run", attributes: "bold" }
+// ABSENT: return

idea/tests/org/jetbrains/jet/completion/KeywordCompletionTestGenerated.java

@@ -287,6 +287,60 @@ public class KeywordCompletionTestGenerated extends AbstractKeywordCompletionTes
         doTest(fileName);
     }
 
+    @TestMetadata("Return1.kt")
+    public void testReturn1() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/keywords/Return1.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("Return2.kt")
+    public void testReturn2() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/keywords/Return2.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("Return3.kt")
+    public void testReturn3() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/keywords/Return3.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("Return4.kt")
+    public void testReturn4() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/keywords/Return4.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("Return5.kt")
+    public void testReturn5() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/keywords/Return5.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("Return6.kt")
+    public void testReturn6() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/keywords/Return6.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("Return7.kt")
+    public void testReturn7() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/keywords/Return7.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("Return8.kt")
+    public void testReturn8() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/keywords/Return8.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("Return9.kt")
+    public void testReturn9() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/keywords/Return9.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("This.kt")
     public void testThis() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/keywords/This.kt");