From 54ee69b6df3632356aeb63255c224316e57044fd Mon Sep 17 00:00:00 2001
From: Dmitry Batkovich <dmitry.batkovich@jetbrains.com>
Date: Thu, 6 Sep 2018 11:46:20 +0300
Subject: [PATCH] functional expression can be folded quick fix should work in
 case of statement lambda

---
 .../afterReturnStatementLambda.java                    |  8 ++++++++
 .../trivialMethodReference/afterStatementLambda.java   |  8 ++++++++
 .../beforeReturnStatementLambda.java                   | 10 ++++++++++
 .../trivialMethodReference/beforeStatementLambda.java  | 10 ++++++++++
 .../FunctionalExpressionCanBeFoldedInspection.java     |  8 ++++++--
 5 files changed, 42 insertions(+), 2 deletions(-)
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/afterReturnStatementLambda.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/afterStatementLambda.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/beforeReturnStatementLambda.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/beforeStatementLambda.java

diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/afterReturnStatementLambda.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/afterReturnStatementLambda.java
new file mode 100644
index 000000000000..bcb8be33642a
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/afterReturnStatementLambda.java
@@ -0,0 +1,8 @@
+// "Replace with qualifier" "true"
+import java.util.function.Function;
+
+class Test {
+  void foo(Function<String, String> function) {
+    Function<String, String> another = function;
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/afterStatementLambda.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/afterStatementLambda.java
new file mode 100644
index 000000000000..5f5f9fcb2c8e
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/afterStatementLambda.java
@@ -0,0 +1,8 @@
+// "Replace with qualifier" "true"
+import java.util.function.Consumer;
+
+class Test {
+  void foo(Consumer<String> consumer) {
+    Consumer<String> another = consumer;
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/beforeReturnStatementLambda.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/beforeReturnStatementLambda.java
new file mode 100644
index 000000000000..7eec771262e6
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/beforeReturnStatementLambda.java
@@ -0,0 +1,10 @@
+// "Replace with qualifier" "true"
+import java.util.function.Function;
+
+class Test {
+  void foo(Function<String, String> function) {
+    Function<String, String> another = s -> {
+      return function.ap<caret>ply(s);
+    };
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/beforeStatementLambda.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/beforeStatementLambda.java
new file mode 100644
index 000000000000..ed874d6ce1dc
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/trivialMethodReference/beforeStatementLambda.java
@@ -0,0 +1,10 @@
+// "Replace with qualifier" "true"
+import java.util.function.Consumer;
+
+class Test {
+  void foo(Consumer<String> consumer) {
+    Consumer<String> another = s -> {
+      consumer.acce<caret>pt(s);
+    };
+  }
+}
\ No newline at end of file
diff --git a/plugins/InspectionGadgets/src/com/intellij/codeInspection/FunctionalExpressionCanBeFoldedInspection.java b/plugins/InspectionGadgets/src/com/intellij/codeInspection/FunctionalExpressionCanBeFoldedInspection.java
index a923839e4a3f..233511b5cecd 100644
--- a/plugins/InspectionGadgets/src/com/intellij/codeInspection/FunctionalExpressionCanBeFoldedInspection.java
+++ b/plugins/InspectionGadgets/src/com/intellij/codeInspection/FunctionalExpressionCanBeFoldedInspection.java
@@ -4,6 +4,7 @@ package com.intellij.codeInspection;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.*;
 import com.intellij.psi.util.MethodSignatureUtil;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.TypeConversionUtil;
 import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
@@ -72,14 +73,17 @@ public class FunctionalExpressionCanBeFoldedInspection extends AbstractBaseJavaL
     @Override
     public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) {
       PsiElement element = descriptor.getPsiElement();
-      final PsiElement parent = element != null ? element.getParent() : null;
+      PsiElement parent = element != null ? element.getParent() : null;
       if (parent instanceof PsiMethodReferenceExpression) {
         final PsiExpression qualifierExpression = ((PsiMethodReferenceExpression)parent).getQualifierExpression();
         if (qualifierExpression != null) {
           parent.replace(qualifierExpression);
         }
       }
-      else if (parent instanceof PsiLambdaExpression) {
+      if (parent instanceof PsiReturnStatement || parent instanceof PsiExpressionStatement) {
+        parent = PsiTreeUtil.getParentOfType(parent, PsiLambdaExpression.class);
+      }
+      if (parent instanceof PsiLambdaExpression) {
         PsiExpression expression = LambdaUtil.extractSingleExpressionFromBody(((PsiLambdaExpression)parent).getBody());
         if (expression instanceof PsiMethodCallExpression) {
           PsiExpression qualifierExpression = ((PsiMethodCallExpression)expression).getMethodExpression().getQualifierExpression();