From 4230df9ede69a00323ef8e0322fb974fa8639eef Mon Sep 17 00:00:00 2001
From: Tagir Valeev <Tagir.Valeev@jetbrains.com>
Date: Fri, 17 Feb 2017 11:11:23 +0700
Subject: [PATCH] StreamToLoop: VoidBlockLambdaFunctionHelper: recreate body if
 invalid (fixes EA-97022 - PIEAE: CompositePsiElement.getContainingFile)

---
 .../codeInspection/streamToLoop/FunctionHelper.java        | 3 +++
 .../quickFix/streamToLoop/afterBlockVoidLambda.java        | 7 +++++++
 .../quickFix/streamToLoop/beforeBlockVoidLambda.java       | 7 +++++++
 3 files changed, 17 insertions(+)

diff --git a/java/java-impl/src/com/intellij/codeInspection/streamToLoop/FunctionHelper.java b/java/java-impl/src/com/intellij/codeInspection/streamToLoop/FunctionHelper.java
index de2c4e3a7544..f356a4c3f80c 100644
--- a/java/java-impl/src/com/intellij/codeInspection/streamToLoop/FunctionHelper.java
+++ b/java/java-impl/src/com/intellij/codeInspection/streamToLoop/FunctionHelper.java
@@ -667,6 +667,9 @@ abstract class FunctionHelper {
 
     void transform(StreamToLoopReplacementContext context, String... argumentValues) {
       super.transform(context, argumentValues);
+      if(!myBody.isValid()) {
+        myBody = ((PsiLambdaExpression)context.createExpression("()->"+myBody.getText())).getBody();
+      }
       List<PsiReturnStatement> returns = getReturns(myBody);
       String continueStatement = "continue;";
       returns.forEach(ret -> ret.replace(context.createStatement(continueStatement)));
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterBlockVoidLambda.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterBlockVoidLambda.java
index 441e7e702618..11bccd6de880 100644
--- a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterBlockVoidLambda.java
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/afterBlockVoidLambda.java
@@ -74,4 +74,11 @@ public class Main {
         }
         System.out.println(count);
     }
+
+    void test(List<String> list) {
+        for (String x : list) {
+            if (x.isEmpty()) continue;
+            System.out.println(x);
+        }
+    }
 }
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeBlockVoidLambda.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeBlockVoidLambda.java
index 9d2d2cedcce5..f63564d4b9d2 100644
--- a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeBlockVoidLambda.java
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamToLoop/beforeBlockVoidLambda.java
@@ -48,4 +48,11 @@ public class Main {
             System.out.println("In flatmap: "+l);
         })).count());
     }
+
+    void test(List<String> list) {
+        list.forEach(x -> {
+            if(x.isEmpty()) return;
+            System.out.println(x);
+        });
+    }
 }
\ No newline at end of file