From 6b5f30f0f2176238a1af55afb253461f2ab3a001 Mon Sep 17 00:00:00 2001
From: peter <peter@jetbrains.com>
Date: Wed, 18 May 2016 14:23:58 +0200
Subject: [PATCH] java completion: prefer return inside then/else parts of if
 (IDEA-156192)

---
 .../completion/PreferByKindWeigher.java       | 30 +++++++++----------
 .../PreferReturnInSingleStatementPlace.java   | 12 ++++++++
 .../NormalCompletionOrderingTest.groovy       |  4 +++
 3 files changed, 31 insertions(+), 15 deletions(-)
 create mode 100644 java/java-tests/testData/codeInsight/completion/normalSorting/PreferReturnInSingleStatementPlace.java

diff --git a/java/java-impl/src/com/intellij/codeInsight/completion/PreferByKindWeigher.java b/java/java-impl/src/com/intellij/codeInsight/completion/PreferByKindWeigher.java
index b1cc29e4e581..84be756a0cc0 100644
--- a/java/java-impl/src/com/intellij/codeInsight/completion/PreferByKindWeigher.java
+++ b/java/java-impl/src/com/intellij/codeInsight/completion/PreferByKindWeigher.java
@@ -171,10 +171,11 @@ public class PreferByKindWeigher extends LookupElementWeigher {
 
     if (object instanceof PsiKeyword) {
       String keyword = ((PsiKeyword)object).getText();
-      if (PsiKeyword.RETURN.equals(keyword) &&
-          isLastStatement(PsiTreeUtil.getParentOfType(myPosition, PsiStatement.class)) &&
-          !isOnTopLevelInVoidMethod(myPosition)) {
-        return MyResult.probableKeyword;
+      if (PsiKeyword.RETURN.equals(keyword)) {
+        PsiStatement parentStatement = PsiTreeUtil.getParentOfType(myPosition, PsiStatement.class);
+        if (isLastStatement(parentStatement) && !isOnTopLevelInVoidMethod(parentStatement)) {
+          return MyResult.probableKeyword;
+        }
       }
       if (PsiKeyword.ELSE.equals(keyword) || PsiKeyword.FINALLY.equals(keyword)) {
         return MyResult.probableKeyword;
@@ -268,17 +269,16 @@ public class PreferByKindWeigher extends LookupElementWeigher {
     return MyResult.normal;
   }
 
-  private static boolean isOnTopLevelInVoidMethod(PsiElement position) {
-    PsiCodeBlock block = PsiTreeUtil.getParentOfType(position, PsiCodeBlock.class);
-    if (block != null) {
-      PsiElement parent = block.getParent();
-      if (parent instanceof PsiMethod) {
-        return ((PsiMethod)parent).isConstructor() || PsiType.VOID.equals(((PsiMethod)parent).getReturnType());
-      }
-      if (parent instanceof PsiLambdaExpression) {
-        PsiMethod method = LambdaUtil.getFunctionalInterfaceMethod(((PsiLambdaExpression)parent).getFunctionalInterfaceType());
-        return method != null && PsiType.VOID.equals(method.getReturnType());
-      }
+  private static boolean isOnTopLevelInVoidMethod(PsiStatement statement) {
+    if (!(statement.getParent() instanceof PsiCodeBlock)) return false;
+
+    PsiElement parent = statement.getParent().getParent();
+    if (parent instanceof PsiMethod) {
+      return ((PsiMethod)parent).isConstructor() || PsiType.VOID.equals(((PsiMethod)parent).getReturnType());
+    }
+    if (parent instanceof PsiLambdaExpression) {
+      PsiMethod method = LambdaUtil.getFunctionalInterfaceMethod(((PsiLambdaExpression)parent).getFunctionalInterfaceType());
+      return method != null && PsiType.VOID.equals(method.getReturnType());
     }
     return false;
   }
diff --git a/java/java-tests/testData/codeInsight/completion/normalSorting/PreferReturnInSingleStatementPlace.java b/java/java-tests/testData/codeInsight/completion/normalSorting/PreferReturnInSingleStatementPlace.java
new file mode 100644
index 000000000000..6da28aa3bb6e
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/completion/normalSorting/PreferReturnInSingleStatementPlace.java
@@ -0,0 +1,12 @@
+import javax.swing.*;
+
+class A extends JPanel {
+    void f() {
+        // "return"
+        if (true) re<caret>
+
+        // "continue"
+        for (int i=0;i<10;i++) {
+        }
+    }
+}
\ No newline at end of file
diff --git a/java/java-tests/testSrc/com/intellij/codeInsight/completion/NormalCompletionOrderingTest.groovy b/java/java-tests/testSrc/com/intellij/codeInsight/completion/NormalCompletionOrderingTest.groovy
index be5f695f72a7..573a79f82ac6 100644
--- a/java/java-tests/testSrc/com/intellij/codeInsight/completion/NormalCompletionOrderingTest.groovy
+++ b/java/java-tests/testSrc/com/intellij/codeInsight/completion/NormalCompletionOrderingTest.groovy
@@ -300,6 +300,10 @@ public class NormalCompletionOrderingTest extends CompletionSortingTestCase {
     checkPreferredItems(0, "return", "rLocal", "rParam", "rMethod");
   }
 
+  public void testPreferReturnInSingleStatementPlace() {
+    checkPreferredItems 0, "return", "registerKeyboardAction"
+  }
+
   public void testPreferModifiers() {
     checkPreferredItems(0, "private", "protected", "public");
   }