From 276b99857a65d83d0c793071c2451a6364e36d2f Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Fri, 26 Dec 2014 19:24:10 +0300
Subject: [PATCH] Insertion of space after "return" depending on return type

---
 .../jet/plugin/completion/CompletionUtils.kt  | 34 +++++++++++++++----
 .../handlers/KotlinKeywordInsertHandler.kt    | 19 ++---------
 2 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/idea/src/org/jetbrains/jet/plugin/completion/CompletionUtils.kt b/idea/src/org/jetbrains/jet/plugin/completion/CompletionUtils.kt
index 8de4f9c4c7e..15c310008f4 100644
--- a/idea/src/org/jetbrains/jet/plugin/completion/CompletionUtils.kt
+++ b/idea/src/org/jetbrains/jet/plugin/completion/CompletionUtils.kt
@@ -63,6 +63,8 @@ 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
+import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns
+import org.jetbrains.jet.plugin.completion.handlers.WithTailInsertHandler
 
 enum class ItemPriority {
     MULTIPLE_ARGUMENTS_ITEM
@@ -278,11 +280,12 @@ private fun functionLiteralLabelAndCall(functionLiteral: JetFunctionLiteral): Pa
 fun returnExpressionItems(bindingContext: BindingContext, position: JetElement): Collection<LookupElement> {
     val result = ArrayList<LookupElement>()
     for (parent in position.parents()) {
-        when (parent) {
-            is JetFunctionLiteral -> {
+        if (parent is JetDeclarationWithBody) {
+            val returnsUnit = returnsUnit(parent, bindingContext)
+            if (parent is JetFunctionLiteral) {
                 val (label, call) = functionLiteralLabelAndCall(parent)
                 if (label != null) {
-                    result.add(createKeywordWithLabelElement("return", label))
+                    result.add(createKeywordWithLabelElement("return", label, addSpace = !returnsUnit))
                 }
 
                 // check if the current function literal is inlined and stop processing outer declarations if it's not
@@ -290,10 +293,9 @@ fun returnExpressionItems(bindingContext: BindingContext, position: JetElement):
                 val target = bindingContext[BindingContext.REFERENCE_TARGET, callee] as? SimpleFunctionDescriptor ?: break // not inlined
                 if (!target.getInlineStrategy().isInline()) break // not inlined
             }
-
-            is JetDeclarationWithBody -> {
+            else {
                 if (parent.hasBlockBody()) {
-                    result.add(createKeywordWithLabelElement("return", null))
+                    result.add(createKeywordWithLabelElement("return", null, addSpace = !returnsUnit))
                 }
                 break
             }
@@ -302,6 +304,26 @@ fun returnExpressionItems(bindingContext: BindingContext, position: JetElement):
     return result
 }
 
+private fun returnsUnit(declaration: JetDeclarationWithBody, bindingContext: BindingContext): Boolean {
+    val callable = bindingContext[BindingContext.DECLARATION_TO_DESCRIPTOR, declaration] as? CallableDescriptor ?: return true
+    val returnType = callable.getReturnType() ?: return true
+    return KotlinBuiltIns.isUnit(returnType)
+}
+
+private fun createKeywordWithLabelElement(keyword: String, label: String?, addSpace: Boolean): LookupElement {
+    val element = createKeywordWithLabelElement(keyword, label)
+    return if (addSpace) {
+        object: LookupElementDecorator<LookupElement>(element) {
+            override fun handleInsert(context: InsertionContext) {
+                WithTailInsertHandler.spaceTail().handleInsert(context, getDelegate())
+            }
+        }
+    }
+    else {
+        element
+    }
+}
+
 private fun createKeywordWithLabelElement(keyword: String, label: String?): LookupElementBuilder {
     var element = LookupElementBuilder.create(KeywordLookupObject, if (label == null) keyword else "$keyword@$label")
     element = element.withPresentableText(keyword)
diff --git a/idea/src/org/jetbrains/jet/plugin/completion/handlers/KotlinKeywordInsertHandler.kt b/idea/src/org/jetbrains/jet/plugin/completion/handlers/KotlinKeywordInsertHandler.kt
index 0d5839ec198..47702896d9d 100644
--- a/idea/src/org/jetbrains/jet/plugin/completion/handlers/KotlinKeywordInsertHandler.kt
+++ b/idea/src/org/jetbrains/jet/plugin/completion/handlers/KotlinKeywordInsertHandler.kt
@@ -35,25 +35,12 @@ public object KotlinKeywordInsertHandler : InsertHandler<LookupElement> {
                                        JetTokens.CONTINUE_KEYWORD.toString())
 
     override fun handleInsert(context: InsertionContext, item: LookupElement) {
-        if (shouldInsertSpaceAfter(item.getLookupString(), context)) {
+        if (shouldInsertSpaceAfter(item.getLookupString())) {
             WithTailInsertHandler.spaceTail().postHandleInsert(context, item)
         }
     }
 
-    private fun shouldInsertSpaceAfter(keyword: String, context: InsertionContext): Boolean {
-        if (keyword in NO_SPACE_AFTER) return false
-
-        if (keyword == JetTokens.RETURN_KEYWORD.toString()) {
-            val element = context.getFile().findElementAt(context.getStartOffset())
-            if (element != null) {
-                val jetFunction = element.getStrictParentOfType<JetFunction>()
-                if (jetFunction != null && (!jetFunction.hasDeclaredReturnType() || JetPsiUtil.isVoidType(jetFunction.getTypeReference()))) {
-                    // No space for void function
-                    return false
-                }
-            }
-        }
-
-        return true
+    private fun shouldInsertSpaceAfter(keyword: String): Boolean {
+        return keyword !in NO_SPACE_AFTER
     }
 }