diff --git a/java/java-analysis-api/src/com/intellij/codeInsight/intention/QuickFixFactory.java b/java/java-analysis-api/src/com/intellij/codeInsight/intention/QuickFixFactory.java index bc5345c863ee..6b4e46676e1a 100644 --- a/java/java-analysis-api/src/com/intellij/codeInsight/intention/QuickFixFactory.java +++ b/java/java-analysis-api/src/com/intellij/codeInsight/intention/QuickFixFactory.java @@ -342,7 +342,7 @@ public abstract class QuickFixFactory { @NotNull PsiExpressionList list); @NotNull - public abstract IntentionAction createAddReturnFix(@NotNull PsiMethod method); + public abstract IntentionAction createAddReturnFix(@NotNull PsiParameterListOwner methodOrLambda); @NotNull public abstract IntentionAction createAddVariableInitializerFix(@NotNull PsiVariable variable); diff --git a/java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/analysis/HighlightControlFlowUtil.java b/java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/analysis/HighlightControlFlowUtil.java index ec230f1e6baa..823ed953dcb5 100644 --- a/java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/analysis/HighlightControlFlowUtil.java +++ b/java/java-analysis-impl/src/com/intellij/codeInsight/daemon/impl/analysis/HighlightControlFlowUtil.java @@ -60,6 +60,9 @@ public class HighlightControlFlowUtil { QuickFixAction.registerQuickFixAction(info, QUICK_FIX_FACTORY.createAddReturnFix(method)); QuickFixAction.registerQuickFixAction(info, QUICK_FIX_FACTORY.createMethodReturnFix(method, PsiType.VOID, true)); } + if (parent instanceof PsiLambdaExpression) { + QuickFixAction.registerQuickFixAction(info, QUICK_FIX_FACTORY.createAddReturnFix((PsiLambdaExpression)parent)); + } return info; } } diff --git a/java/java-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/AddReturnFix.java b/java/java-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/AddReturnFix.java index f8df85fd4b38..b60060eedeb3 100644 --- a/java/java-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/AddReturnFix.java +++ b/java/java-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/AddReturnFix.java @@ -26,7 +26,9 @@ import com.intellij.psi.util.InheritanceUtil; import com.intellij.psi.util.PsiTypesUtil; import com.intellij.psi.util.PsiUtil; import com.intellij.psi.util.TypeConversionUtil; +import com.intellij.util.ObjectUtils; import com.intellij.util.containers.ContainerUtil; +import com.siyeh.ig.psiutils.CommentTracker; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -34,10 +36,11 @@ import java.util.ArrayList; import java.util.List; public class AddReturnFix implements IntentionAction { - private final PsiMethod myMethod; + private final PsiParameterListOwner myMethod; - public AddReturnFix(@NotNull PsiMethod method) { - myMethod = method; + + public AddReturnFix(@NotNull PsiParameterListOwner methodOrLambda) { + myMethod = methodOrLambda; } @Override @@ -56,9 +59,8 @@ public class AddReturnFix implements IntentionAction { public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { return myMethod.isValid() && BaseIntentionAction.canModify(myMethod) && - myMethod.getBody() != null && - myMethod.getBody().getRBrace() != null - ; + myMethod.getBody() instanceof PsiCodeBlock && + ((PsiCodeBlock)myMethod.getBody()).getRBrace() != null; } @NotNull @@ -69,19 +71,33 @@ public class AddReturnFix implements IntentionAction { @Override public void invoke(@NotNull Project project, Editor editor, PsiFile file) { + if (invokeSingleExpressionLambdaFix()) { + return; + } String value = suggestReturnValue(); PsiElementFactory factory = JavaPsiFacade.getElementFactory(myMethod.getProject()); PsiReturnStatement returnStatement = (PsiReturnStatement) factory.createStatementFromText("return " + value+";", myMethod); - PsiCodeBlock body = myMethod.getBody(); + PsiCodeBlock body = ObjectUtils.tryCast(myMethod.getBody(), PsiCodeBlock.class); + assert body != null; returnStatement = (PsiReturnStatement) body.addBefore(returnStatement, body.getRBrace()); MethodReturnTypeFix.selectReturnValueInEditor(returnStatement, editor); } private String suggestReturnValue() { - PsiType type = myMethod.getReturnType(); + PsiType type; + if (myMethod instanceof PsiMethod) { + type = ((PsiMethod)myMethod).getReturnType(); + } + else if (myMethod instanceof PsiLambdaExpression) { + type = LambdaUtil.getFunctionalInterfaceReturnType((PsiLambdaExpression)myMethod); + } + else { + return PsiKeyword.NULL; // normally shouldn't happen + } + // first try to find suitable local variable - PsiVariable[] variables = getDeclaredVariables(myMethod); + List variables = getDeclaredVariables(myMethod); for (PsiVariable variable : variables) { PsiType varType = variable.getType(); if (varType.equals(type)) { @@ -127,20 +143,51 @@ public class AddReturnFix implements IntentionAction { } @NotNull - private static PsiVariable[] getDeclaredVariables(PsiMethod method) { + private static List getDeclaredVariables(PsiParameterListOwner method) { List variables = new ArrayList<>(); - PsiStatement[] statements = method.getBody().getStatements(); - for (PsiStatement statement : statements) { - if (statement instanceof PsiDeclarationStatement) { - PsiElement[] declaredElements = ((PsiDeclarationStatement)statement).getDeclaredElements(); - for (PsiElement declaredElement : declaredElements) { - if (declaredElement instanceof PsiLocalVariable) variables.add((PsiVariable)declaredElement); + PsiCodeBlock body = ObjectUtils.tryCast(method.getBody(), PsiCodeBlock.class); + if (body != null) { + PsiStatement[] statements = body.getStatements(); + for (PsiStatement statement : statements) { + if (statement instanceof PsiDeclarationStatement) { + PsiElement[] declaredElements = ((PsiDeclarationStatement)statement).getDeclaredElements(); + for (PsiElement declaredElement : declaredElements) { + if (declaredElement instanceof PsiLocalVariable) variables.add((PsiVariable)declaredElement); + } } } } PsiParameter[] parameters = method.getParameterList().getParameters(); ContainerUtil.addAll(variables, parameters); - return variables.toArray(new PsiVariable[0]); + return variables; + } + + private boolean invokeSingleExpressionLambdaFix() { + if (myMethod instanceof PsiLambdaExpression) { + PsiType returnType = LambdaUtil.getFunctionalInterfaceReturnType((PsiLambdaExpression)myMethod); + if (returnType != null) { + PsiCodeBlock body = ObjectUtils.tryCast(myMethod.getBody(), PsiCodeBlock.class); + if (body != null) { + PsiStatement[] statements = body.getStatements(); + if (statements.length == 1 && statements[0] instanceof PsiExpressionStatement) { + PsiExpressionStatement expressionStatement = (PsiExpressionStatement)statements[0]; + PsiExpression expression = expressionStatement.getExpression(); + PsiType expressionType = expression.getType(); + + if (expressionType != null && returnType.isAssignableFrom(expressionType)) { + PsiElementFactory factory = JavaPsiFacade.getElementFactory(myMethod.getProject()); + PsiReturnStatement returnStatement = (PsiReturnStatement)factory.createStatementFromText("return 0;", myMethod); + ObjectUtils.notNull(returnStatement.getReturnValue()).replace(expression); + CommentTracker tracker = new CommentTracker(); + tracker.markUnchanged(expression); + tracker.replaceAndRestoreComments(expressionStatement, returnStatement); + return true; + } + } + } + } + } + return false; } @Override diff --git a/java/java-impl/src/com/intellij/codeInsight/intention/impl/config/QuickFixFactoryImpl.java b/java/java-impl/src/com/intellij/codeInsight/intention/impl/config/QuickFixFactoryImpl.java index c0d18193a905..49610e6a7447 100644 --- a/java/java-impl/src/com/intellij/codeInsight/intention/impl/config/QuickFixFactoryImpl.java +++ b/java/java-impl/src/com/intellij/codeInsight/intention/impl/config/QuickFixFactoryImpl.java @@ -584,8 +584,8 @@ public class QuickFixFactoryImpl extends QuickFixFactory { @NotNull @Override - public IntentionAction createAddReturnFix(@NotNull PsiMethod method) { - return new AddReturnFix(method); + public IntentionAction createAddReturnFix(@NotNull PsiParameterListOwner methodOrLambda) { + return new AddReturnFix(methodOrLambda); } @NotNull diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda1.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda1.java new file mode 100644 index 000000000000..4749adc0062d --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda1.java @@ -0,0 +1,14 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + return 0; + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda2.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda2.java new file mode 100644 index 000000000000..4f284b99ab44 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda2.java @@ -0,0 +1,14 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + return 1; + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda3.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda3.java new file mode 100644 index 000000000000..d38b1e771d74 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda3.java @@ -0,0 +1,14 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + return Math.max(1, 2); + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda4.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda4.java new file mode 100644 index 000000000000..d38b1e771d74 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda4.java @@ -0,0 +1,14 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + return Math.max(1, 2); + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda5.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda5.java new file mode 100644 index 000000000000..58647010fdc3 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda5.java @@ -0,0 +1,14 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + return Math.max(1, 2); //comment + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda6.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda6.java new file mode 100644 index 000000000000..7a9cf141acda --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/afterLambda6.java @@ -0,0 +1,14 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + return Math.max(1,/*comment*/ 2); + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda1.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda1.java new file mode 100644 index 000000000000..52b3cba0583c --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda1.java @@ -0,0 +1,13 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda2.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda2.java new file mode 100644 index 000000000000..02e5c7edccfa --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda2.java @@ -0,0 +1,14 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + 1 + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda3.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda3.java new file mode 100644 index 000000000000..0729945faebc --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda3.java @@ -0,0 +1,14 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + Math.max(1, 2) + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda4.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda4.java new file mode 100644 index 000000000000..442962d1ce13 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda4.java @@ -0,0 +1,14 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + Math.max(1, 2); + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda5.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda5.java new file mode 100644 index 000000000000..bf1fb4854a39 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda5.java @@ -0,0 +1,14 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + Math.max(1, 2) //comment + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda6.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda6.java new file mode 100644 index 000000000000..8faccc37f300 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addReturn/beforeLambda6.java @@ -0,0 +1,14 @@ +// "Add 'return' statement" "true" +class C { + void foo() { + bar(() -> { + Math.max(1,/*comment*/ 2); + }); + } + + void bar(I i) {} + + interface I { + int f(); + } +} \ No newline at end of file