From 1279ab20d81cc820a10db72556810b14d15785d0 Mon Sep 17 00:00:00 2001 From: Dmitry Batkovich Date: Mon, 29 May 2017 19:15:24 +0300 Subject: [PATCH] method chain completion: make it works inside return statement --- .../MethodChainCompletionContributor.java | 29 +++++++++++++++---- .../testReturnStatement/TestCompletion.java | 15 ++++++++++ .../testReturnStatement/TestIndex.java | 20 +++++++++++++ .../MethodChainsCompletionTest.java | 5 ++++ 4 files changed, 64 insertions(+), 5 deletions(-) create mode 100644 java/java-tests/testData/codeInsight/completion/methodChains/testReturnStatement/TestCompletion.java create mode 100644 java/java-tests/testData/codeInsight/completion/methodChains/testReturnStatement/TestIndex.java diff --git a/java/compiler/impl/src/com/intellij/compiler/chainsSearch/completion/MethodChainCompletionContributor.java b/java/compiler/impl/src/com/intellij/compiler/chainsSearch/completion/MethodChainCompletionContributor.java index e5c722ecab86..f7babe6e06d4 100644 --- a/java/compiler/impl/src/com/intellij/compiler/chainsSearch/completion/MethodChainCompletionContributor.java +++ b/java/compiler/impl/src/com/intellij/compiler/chainsSearch/completion/MethodChainCompletionContributor.java @@ -31,8 +31,8 @@ import com.intellij.openapi.diagnostic.Logger; import com.intellij.openapi.util.registry.Registry; import com.intellij.patterns.ElementPattern; import com.intellij.psi.*; -import com.intellij.psi.impl.source.tree.java.PsiMethodCallExpressionImpl; import com.intellij.psi.util.PsiTreeUtil; +import com.intellij.psi.util.PsiTypesUtil; import com.intellij.util.ProcessingContext; import gnu.trove.THashSet; import org.jetbrains.annotations.NotNull; @@ -52,7 +52,7 @@ public class MethodChainCompletionContributor extends CompletionContributor { private static final boolean UNIT_TEST_MODE = ApplicationManager.getApplication().isUnitTestMode(); public MethodChainCompletionContributor() { - ElementPattern pattern = or(patternForMethodCallArgument(), patternForVariableAssignment()); + ElementPattern pattern = or(patternForMethodCallArgument(), patternForVariableAssignment(), patternForReturnExpression()); extend(CompletionType.SMART, pattern, new CompletionProvider() { @Override protected void addCompletions(@NotNull CompletionParameters parameters, @@ -109,9 +109,16 @@ public class MethodChainCompletionContributor extends CompletionContributor { @Nullable private static ChainCompletionContext extractContext(CompletionParameters parameters) { - PsiElement parent = PsiTreeUtil.getParentOfType(parameters.getPosition(), PsiAssignmentExpression.class, PsiLocalVariable.class, PsiMethodCallExpression.class); + PsiElement parent = PsiTreeUtil.getParentOfType(parameters.getPosition(), + PsiAssignmentExpression.class, + PsiLocalVariable.class, + PsiMethodCallExpression.class, + PsiReturnStatement.class); LOG.assertTrue(parent != null, "A completion position should match to a pattern"); + if (parent instanceof PsiReturnStatement) { + return extractContextFromReturn((PsiReturnStatement)parent, parameters); + } if (parent instanceof PsiAssignmentExpression) { return extractContextFromAssignment((PsiAssignmentExpression)parent, parameters); } @@ -127,11 +134,19 @@ public class MethodChainCompletionContributor extends CompletionContributor { PsiParameter[] methodParameters = method.getParameterList().getParameters(); if (exprPosition < methodParameters.length) { PsiParameter methodParameter = methodParameters[exprPosition]; - return ChainCompletionContext.createContext(methodParameter.getType(), PsiTreeUtil.getParentOfType(expression, PsiDeclarationStatement.class), suggestIterators(parameters)); + return ChainCompletionContext.createContext(methodParameter.getType(), PsiTreeUtil.getParentOfType(expression, PsiMethodCallExpression.class), suggestIterators(parameters)); } return null; } + @Nullable + private static ChainCompletionContext extractContextFromReturn(PsiReturnStatement returnStatement, + CompletionParameters parameters) { + PsiType type = PsiTypesUtil.getMethodReturnType(returnStatement); + if (type == null) return null; + return ChainCompletionContext.createContext(type, returnStatement, suggestIterators(parameters)); + } + @Nullable private static ChainCompletionContext extractContextFromVariable(PsiLocalVariable localVariable, CompletionParameters parameters) { @@ -163,7 +178,11 @@ public class MethodChainCompletionContributor extends CompletionContributor { @NotNull private static ElementPattern patternForMethodCallArgument() { - return psiElement().withSuperParent(3, PsiMethodCallExpressionImpl.class); + return psiElement().withSuperParent(3, PsiMethodCallExpression.class); + } + + private static ElementPattern patternForReturnExpression() { + return psiElement().withSuperParent(2, PsiReturnStatement.class); } private static boolean suggestIterators(@NotNull CompletionParameters parameters) { diff --git a/java/java-tests/testData/codeInsight/completion/methodChains/testReturnStatement/TestCompletion.java b/java/java-tests/testData/codeInsight/completion/methodChains/testReturnStatement/TestCompletion.java new file mode 100644 index 000000000000..de11826dc6f0 --- /dev/null +++ b/java/java-tests/testData/codeInsight/completion/methodChains/testReturnStatement/TestCompletion.java @@ -0,0 +1,15 @@ +class PsiClass { +} + +class PsiElementFactory { + public PsiClass createClass() { + return null; + } +} + +public class TestCompletion { + + public PsiClass method(PsiElementFactory f) { + return + } +} diff --git a/java/java-tests/testData/codeInsight/completion/methodChains/testReturnStatement/TestIndex.java b/java/java-tests/testData/codeInsight/completion/methodChains/testReturnStatement/TestIndex.java new file mode 100644 index 000000000000..4c670a3359b2 --- /dev/null +++ b/java/java-tests/testData/codeInsight/completion/methodChains/testReturnStatement/TestIndex.java @@ -0,0 +1,20 @@ +public class TestIndex { + + public void statMethod(PsiElementFactory f) { + f.createClass(); + f.createClass(); + f.createClass(); + f.createClass(); + f.createClass(); + f.createClass(); + } +} + +class PsiClass { +} + +class PsiElementFactory { + public PsiClass createClass() { + return null; + } +} \ No newline at end of file diff --git a/java/java-tests/testSrc/com/intellij/java/codeInsight/completion/MethodChainsCompletionTest.java b/java/java-tests/testSrc/com/intellij/java/codeInsight/completion/MethodChainsCompletionTest.java index 2cd9978a2b2e..dfe97f1d6eab 100644 --- a/java/java-tests/testSrc/com/intellij/java/codeInsight/completion/MethodChainsCompletionTest.java +++ b/java/java-tests/testSrc/com/intellij/java/codeInsight/completion/MethodChainsCompletionTest.java @@ -221,6 +221,11 @@ public class MethodChainsCompletionTest extends AbstractCompilerAwareTest { assertEquals("c.getMethods", element.getLookupString()); } + public void testReturnStatement() { + JavaRelevantChainLookupElement element = assertOneElement(doCompletion()); + assertEquals("f.createClass", element.getLookupString()); + } + public void assertAdvisorLookupElementEquals(String lookupText, int unreachableParametersCount, int chainSize,