KT-1238 Completion of 'break' keyword: why does it add space char at the end?

KT-1237 Auto-completion of "return" keyword should not insert space after it when in a method returning Unit

 #KT-1237 fixed
 #KT-1238 fixed

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiUtil.java

@@ -299,4 +299,12 @@ public class JetPsiUtil {
         }
         return null;
     }
+
+    public static boolean isVoidType(@Nullable JetTypeReference typeReference) {
+        if (typeReference == null) {
+            return false;
+        }
+
+        return "Unit".equals(typeReference.getText());
+    }
 }

idea/src/org/jetbrains/jet/plugin/completion/handlers/JetKeywordInsertHandler.java

@@ -21,6 +21,11 @@ import com.intellij.codeInsight.TailType;
 import com.intellij.codeInsight.completion.InsertHandler;
 import com.intellij.codeInsight.completion.InsertionContext;
 import com.intellij.codeInsight.lookup.LookupElement;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
+import org.jetbrains.jet.lang.psi.JetFunction;
+import org.jetbrains.jet.lang.psi.JetPsiUtil;
+import org.jetbrains.jet.lexer.JetTokens;
 
 import java.util.Set;
 
@@ -29,17 +34,40 @@ import java.util.Set;
  */
 public class JetKeywordInsertHandler implements InsertHandler<LookupElement> {
 
-    private final static Set<String> NO_SPACE_AFTER = Sets.newHashSet("this", "super", "This", "true", "false", "null");
+    private final static Set<String> NO_SPACE_AFTER = Sets.newHashSet(
+            JetTokens.THIS_KEYWORD.toString(),
+            JetTokens.SUPER_KEYWORD.toString(),
+            JetTokens.CAPITALIZED_THIS_KEYWORD.toString(),
+            JetTokens.THIS_KEYWORD.toString(),
+            JetTokens.FALSE_KEYWORD.toString(),
+            JetTokens.NULL_KEYWORD.toString(),
+            JetTokens.BREAK_KEYWORD.toString(),
+            JetTokens.CONTINUE_KEYWORD.toString());
 
     @Override
     public void handleInsert(InsertionContext context, LookupElement item) {
         String keyword = item.getLookupString();
 
-        // Add space after keyword
-        if (!NO_SPACE_AFTER.contains(keyword)) {
-            context.setAddCompletionChar(false);
-            final TailType tailType = TailType.SPACE;
-            tailType.processTail(context.getEditor(), context.getTailOffset());
+        if (NO_SPACE_AFTER.contains(keyword)) {
+            return;
         }
+
+        if (keyword.equals(JetTokens.RETURN_KEYWORD.toString())) {
+            PsiElement element = context.getFile().findElementAt(context.getStartOffset());
+            if (element != null) {
+                JetFunction jetFunction = PsiTreeUtil.getParentOfType(element, JetFunction.class);
+                if (jetFunction != null) {
+                    if (!jetFunction.hasDeclaredReturnType() || JetPsiUtil.isVoidType(jetFunction.getReturnTypeRef())) {
+                        // No space for void function
+                        return;
+                    }
+                }
+            }
+        }
+
+        // Add space after keyword
+        context.setAddCompletionChar(false);
+        final TailType tailType = TailType.SPACE;
+        tailType.processTail(context.getEditor(), context.getTailOffset());
     }
 }

idea/testData/completion/handlers/keywords/Break.kt

@@ -0,0 +1,3 @@
+fun t() {
+    brea<caret>
+}

idea/testData/completion/handlers/keywords/Break_after.kt

@@ -0,0 +1,3 @@
+fun t() {
+    break<caret>
+}

idea/testData/completion/handlers/keywords/ReturnInEmptyType.kt

@@ -0,0 +1,3 @@
+fun test(a: Int) {
+    retur<caret>
+}

idea/testData/completion/handlers/keywords/ReturnInEmptyType_after.kt

@@ -0,0 +1,3 @@
+fun test(a: Int) {
+    return<caret>
+}

idea/testData/completion/handlers/keywords/ReturnInProperty.kt

@@ -0,0 +1,6 @@
+class Test {
+    val test : String
+        get() {
+            retur<caret>
+        }
+}

idea/testData/completion/handlers/keywords/ReturnInProperty_after.kt

@@ -0,0 +1,6 @@
+class Test {
+    val test : String
+        get() {
+            return <caret>
+        }
+}

idea/testData/completion/handlers/keywords/ReturnInTypeFunction.kt

@@ -0,0 +1,5 @@
+class Test {
+    fun someMethod() : Int {
+        retur<caret>
+    }
+}

idea/testData/completion/handlers/keywords/ReturnInTypeFunction_after.kt

@@ -0,0 +1,5 @@
+class Test {
+    fun someMethod() : Int {
+        return <caret>
+    }
+}

idea/testData/completion/handlers/keywords/ReturnInUnit.kt

@@ -0,0 +1,5 @@
+object OtherTest {
+    fun test() : Unit {
+        retur<caret>
+    }
+}

idea/testData/completion/handlers/keywords/ReturnInUnit_after.kt

@@ -0,0 +1,5 @@
+object OtherTest {
+    fun test() : Unit {
+        return<caret>
+    }
+}

idea/tests/org/jetbrains/jet/completion/handlers/KeywordsHandlerTest.java

@@ -39,9 +39,35 @@ public class KeywordsHandlerTest extends LightCompletionTestCase {
         checkResultByText("fun test() { null<caret> }");
     }
 
+    public void testBreak() {
+        doFileTest();
+    }
+
+    public void testReturnInEmptyType() {
+        doFileTest();
+    }
+
+    public void testReturnInProperty() {
+        doFileTest();
+    }
+
+    public void testReturnInTypeFunction() {
+        doFileTest();
+    }
+
+    public void testReturnInUnit() {
+        doFileTest();
+    }
+
+    protected void doFileTest() {
+        configureByFile(getTestName(false) + ".kt");
+        complete();
+        checkResultByFile(getTestName(false) + "_after.kt");
+    }
+
     @Override
     protected String getTestDataPath() {
-        return new File(PluginTestCaseBase.getTestDataPathBase(), "/completion/handlers/").getPath() +
+        return new File(PluginTestCaseBase.getTestDataPathBase(), "/completion/handlers/keywords").getPath() +
                File.separator;
     }
 }