From a6e35cfd57a15fa66adff75643d26f9f1a15880d Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Mon, 27 Oct 2014 12:14:28 +0100
Subject: [PATCH] calculate void/value lambda compatibility based on valid,
 non-modified lambda expression

---
 .../refactoring/util/RefactoringUtil.java     | 38 +++++++++++--------
 ...utInLambdaBodyVoidValueConflict.after.java | 20 ++++++++++
 .../PutInLambdaBodyVoidValueConflict.java     | 17 +++++++++
 .../refactoring/IntroduceVariableTest.java    |  4 ++
 4 files changed, 63 insertions(+), 16 deletions(-)
 create mode 100644 java/java-tests/testData/refactoring/introduceVariable/PutInLambdaBodyVoidValueConflict.after.java
 create mode 100644 java/java-tests/testData/refactoring/introduceVariable/PutInLambdaBodyVoidValueConflict.java

diff --git a/java/java-impl/src/com/intellij/refactoring/util/RefactoringUtil.java b/java/java-impl/src/com/intellij/refactoring/util/RefactoringUtil.java
index 004599135fee..d181d3c49351 100644
--- a/java/java-impl/src/com/intellij/refactoring/util/RefactoringUtil.java
+++ b/java/java-impl/src/com/intellij/refactoring/util/RefactoringUtil.java
@@ -870,32 +870,38 @@ public class RefactoringUtil {
       JavaCodeStyleManager.getInstance(declaration.getProject()).shortenClassReferences(declaration);
       if (loopBodyCopy != null) codeBlock.add(loopBodyCopy);
     } else if (container instanceof PsiLambdaExpression) {
-      final PsiLambdaExpression lambdaExpression = (PsiLambdaExpression)container;
-      final PsiElement lambdaExpressionBody = lambdaExpression.getBody();
-      LOG.assertTrue(lambdaExpressionBody != null);
+      PsiLambdaExpression lambdaExpression = (PsiLambdaExpression)container;
+      final PsiElement invalidBody = lambdaExpression.getBody();
+      if (invalidBody == null) return declaration;
 
       final PsiLambdaExpression expressionFromText = (PsiLambdaExpression)elementFactory
-        .createExpressionFromText(lambdaExpression.getParameterList().getText() + " -> {}", lambdaExpression);
+        .createExpressionFromText(lambdaExpression.getParameterList().getText() + " -> {}", lambdaExpression.getParent());
       PsiCodeBlock newBody = (PsiCodeBlock)expressionFromText.getBody();
       LOG.assertTrue(newBody != null);
       newBody.add(declaration);
 
-      PsiStatement lastBodyStatement = elementFactory.createStatementFromText("a;", lambdaExpression);
-      ((PsiExpressionStatement)lastBodyStatement).getExpression().replace(lambdaExpressionBody);
+      lambdaExpression =
+        (PsiLambdaExpression)lambdaExpression.replace(elementFactory.createExpressionFromText(
+          lambdaExpression.getParameterList().getText() + " -> " + invalidBody.getText(), lambdaExpression));
+
+      final PsiElement lambdaExpressionBody = lambdaExpression.getBody();
+      LOG.assertTrue(lambdaExpressionBody != null);
+      final PsiStatement lastBodyStatement;
+      if (LambdaUtil.getFunctionalInterfaceReturnType(lambdaExpression) == PsiType.VOID) {
+        lastBodyStatement = elementFactory.createStatementFromText("a;", lambdaExpression);
+        ((PsiExpressionStatement)lastBodyStatement).getExpression().replace(lambdaExpressionBody);
+      }
+      else {
+        lastBodyStatement = elementFactory.createStatementFromText("return a;", lambdaExpression);
+        final PsiExpression returnValue = ((PsiReturnStatement)lastBodyStatement).getReturnValue();
+        LOG.assertTrue(returnValue != null);
+        returnValue.replace(lambdaExpressionBody);
+      }
       newBody.add(lastBodyStatement);
 
-      PsiLambdaExpression copy = (PsiLambdaExpression)lambdaExpression.replace(expressionFromText);
+      final PsiLambdaExpression copy = (PsiLambdaExpression)lambdaExpression.replace(expressionFromText);
       newBody = (PsiCodeBlock)copy.getBody();
       LOG.assertTrue(newBody != null);
-
-      if (LambdaUtil.getFunctionalInterfaceReturnType(copy) != PsiType.VOID) {
-        PsiExpressionStatement lastStatement = (PsiExpressionStatement)newBody.getStatements()[1];
-        PsiReturnStatement returnStatement = (PsiReturnStatement)elementFactory.createStatementFromText("return a;", copy);
-        final PsiExpression returnValue = returnStatement.getReturnValue();
-        LOG.assertTrue(returnValue != null);
-        returnValue.replace(lastStatement.getExpression());
-        lastStatement.replace(returnStatement);
-      }
       declaration = newBody.getStatements()[0];
       declaration = (PsiStatement)JavaCodeStyleManager.getInstance(declaration.getProject()).shortenClassReferences(declaration);
     }
diff --git a/java/java-tests/testData/refactoring/introduceVariable/PutInLambdaBodyVoidValueConflict.after.java b/java/java-tests/testData/refactoring/introduceVariable/PutInLambdaBodyVoidValueConflict.after.java
new file mode 100644
index 000000000000..6315c2590032
--- /dev/null
+++ b/java/java-tests/testData/refactoring/introduceVariable/PutInLambdaBodyVoidValueConflict.after.java
@@ -0,0 +1,20 @@
+class Foo {
+    interface I<T> {
+        T m(Object op);
+    }
+
+    interface J<T> {
+        void m(T o);
+    }
+
+
+    void f(J r) {}
+    void f(I<String> r) {}
+
+    {
+        f((a) -> {
+            int c = 1;
+            return c;
+        });
+    }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/refactoring/introduceVariable/PutInLambdaBodyVoidValueConflict.java b/java/java-tests/testData/refactoring/introduceVariable/PutInLambdaBodyVoidValueConflict.java
new file mode 100644
index 000000000000..bdea215a08a4
--- /dev/null
+++ b/java/java-tests/testData/refactoring/introduceVariable/PutInLambdaBodyVoidValueConflict.java
@@ -0,0 +1,17 @@
+class Foo {
+    interface I<T> {
+        T m(Object op);
+    }
+
+    interface J<T> {
+        void m(T o);
+    }
+
+
+    void f(J r) {}
+    void f(I<String> r) {}
+
+    {
+        f((a) -> <selection>1</selection>);
+    }
+}
\ No newline at end of file
diff --git a/java/java-tests/testSrc/com/intellij/refactoring/IntroduceVariableTest.java b/java/java-tests/testSrc/com/intellij/refactoring/IntroduceVariableTest.java
index 29f792330bbf..2e2a6cc488a3 100644
--- a/java/java-tests/testSrc/com/intellij/refactoring/IntroduceVariableTest.java
+++ b/java/java-tests/testSrc/com/intellij/refactoring/IntroduceVariableTest.java
@@ -441,6 +441,10 @@ public class IntroduceVariableTest extends LightCodeInsightTestCase {
     doTest(new MockIntroduceVariableHandler("c", false, false, false, "int"));
   }
 
+  public void testPutInLambdaBodyVoidValueConflict() {
+    doTest(new MockIntroduceVariableHandler("c", false, false, false, "int"));
+  }
+
   public void testNormalizeDeclarations() {
     doTest(new MockIntroduceVariableHandler("i3", false, false, false, "int"));
   }