wrap return value: expand method reference to lambda

java/java-impl/src/com/intellij/refactoring/wrapreturnvalue/WrapReturnValueProcessor.java

@@ -136,6 +136,9 @@ public class WrapReturnValueProcessor extends FixableUsagesRefactoringProcessor
       if (parent instanceof PsiCallExpression) {
         usages.add(new UnwrapCall((PsiCallExpression)parent, myUnwrapMethodName));
       }
+      else if (referenceElement instanceof PsiMethodReferenceExpression) {
+        usages.add(new UnwrapCall((PsiMethodReferenceExpression)referenceElement, myUnwrapMethodName));
+      }
     }
     final String returnType = calculateReturnTypeString();
     usages.add(new ChangeReturnType(psiMethod, returnType));

java/java-impl/src/com/intellij/refactoring/wrapreturnvalue/usageInfo/UnwrapCall.java

@@ -15,25 +15,33 @@
  */
 package com.intellij.refactoring.wrapreturnvalue.usageInfo;
 
-import com.intellij.psi.PsiCallExpression;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiExpression;
+import com.intellij.psi.PsiMethodReferenceExpression;
 import com.intellij.refactoring.psi.MutationUtils;
 import com.intellij.refactoring.util.FixableUsageInfo;
+import com.intellij.refactoring.util.LambdaRefactoringUtil;
 import com.intellij.util.IncorrectOperationException;
 import org.jetbrains.annotations.NotNull;
 
 public class UnwrapCall extends FixableUsageInfo {
-  private final PsiCallExpression myCall;
   private final String myUnwrapMethod;
 
-  public UnwrapCall(@NotNull PsiCallExpression call, @NotNull String unwrapMethod) {
+  public UnwrapCall(@NotNull PsiExpression call, @NotNull String unwrapMethod) {
     super(call);
-    myCall = call;
     myUnwrapMethod = unwrapMethod;
   }
 
   @Override
   public void fixUsage() throws IncorrectOperationException {
-    String newExpression = myCall.getText() + '.' + myUnwrapMethod + "()";
-    MutationUtils.replaceExpression(newExpression, myCall);
+    PsiElement element = getElement();
+    if (!(element instanceof PsiExpression)) return;
+    if (element instanceof PsiMethodReferenceExpression) {
+      PsiExpression expression = LambdaRefactoringUtil.convertToMethodCallInLambdaBody((PsiMethodReferenceExpression)element);
+      if (expression == null) return;
+      element = expression;
+    }
+    String newExpression = element.getText() + '.' + myUnwrapMethod + "()";
+    MutationUtils.replaceExpression(newExpression, (PsiExpression)element);
   }
 }

java/java-tests/testData/refactoring/wrapReturnValue/openMethodReference/after/Test.java

@@ -0,0 +1,25 @@
+class Test {
+  interface F {
+    String get(Test t);
+  }
+
+    Wrapper foo() {
+    return new Wrapper("");
+  }
+  
+  {
+    F f = test -> test.foo().getValue();
+  }
+
+    public class Wrapper {
+        private final String value;
+
+        public Wrapper(String value) {
+            this.value = value;
+        }
+
+        public String getValue() {
+            return value;
+        }
+    }
+}

java/java-tests/testData/refactoring/wrapReturnValue/openMethodReference/before/Test.java

@@ -0,0 +1,13 @@
+class Test {
+  interface F {
+    String get(Test t);
+  }
+
+  String foo() {
+    return "";
+  }
+  
+  {
+    F f = Test::foo;
+  }
+}

java/java-tests/testSrc/com/intellij/java/refactoring/WrapReturnValueTest.java

@@ -55,6 +55,7 @@ public class WrapReturnValueTest extends MultiFileTestCase {
   public void testInferFieldTypeArg() { doTest(true, null, false); }
   public void testWrongFieldType() { doTest(true, "Existing class does not have appropriate constructor", false); }
   public void testStaticMethodInnerClass() { doTest(false, null, true); }
+  public void testOpenMethodReference() { doTest(false, null, true); }
   public void testRawReturnType() { doTest(true, "Existing class does not have appropriate constructor"); }
   public void testReturnInsideLambda() { doTest(false, null, true); }
   public void testTypeAnnotations() { doTest(false); }