IDEA-56460 'Convert to for-in' removes too much code

plugins/groovy/src/org/jetbrains/plugins/groovy/intentions/closure/EachToForIntention.java

@@ -69,12 +69,13 @@ public class EachToForIntention extends Intention {
     StringBuilder builder = new StringBuilder();
     builder.append("for (").append(var).append(" in ").append(qualifier.getText()).append(") {\n");
     String text = block.getText();
-    int index = text.indexOf("->");
-    if (index == -1) {
-      index = 1;
+    final PsiElement blockArrow = block.getArrow();
+    int index;
+    if (blockArrow != null) {
+      index = blockArrow.getStartOffsetInParent() + blockArrow.getTextLength();
     }
     else {
-      index += 2;
+      index = 1;
     }
     while (index < text.length() && Character.isWhitespace(text.charAt(index))) index++;
     text = text.substring(index, text.length() - 1);

plugins/groovy/src/org/jetbrains/plugins/groovy/lang/psi/api/statements/blocks/GrClosableBlock.java

@@ -16,6 +16,7 @@
 
 package org.jetbrains.plugins.groovy.lang.psi.api.statements.blocks;
 
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiParameter;
 import com.intellij.psi.PsiType;
 import org.jetbrains.annotations.Nullable;
@@ -44,4 +45,7 @@ public interface GrClosableBlock extends GrExpression, GrCodeBlock, GrParameters
   PsiType getReturnType();
 
   PsiParameter[] getAllParameters();
+
+  @Nullable
+  PsiElement getArrow();
 }

plugins/groovy/src/org/jetbrains/plugins/groovy/lang/psi/impl/statements/blocks/GrClosableBlockImpl.java

@@ -24,7 +24,6 @@ import com.intellij.util.Function;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.plugins.groovy.lang.lexer.GroovyTokenTypes;
-import org.jetbrains.plugins.groovy.lang.parser.GroovyElementTypes;
 import org.jetbrains.plugins.groovy.lang.psi.GroovyElementVisitor;
 import org.jetbrains.plugins.groovy.lang.psi.GroovyFile;
 import org.jetbrains.plugins.groovy.lang.psi.GroovyPsiElement;
@@ -114,13 +113,19 @@ public class GrClosableBlockImpl extends GrBlockImpl implements GrClosableBlock
     return new PsiParameter[]{getSyntheticItParameter()};
   }
 
+  @Override
+  @Nullable
+  public PsiElement getArrow() {
+    return findChildByType(GroovyTokenTypes.mCLOSABLE_BLOCK_OP);
+  }
+
   public GrParameterListImpl getParameterList() {
     return findChildByClass(GrParameterListImpl.class);
   }
 
   public void addParameter(GrParameter parameter) {
     GrParameterList parameterList = getParameterList();
-    if (findChildByType(GroovyTokenTypes.mCLOSABLE_BLOCK_OP) == null) {
+    if (getArrow() == null) {
       ASTNode next = parameterList.getNode().getTreeNext();
       getNode().addLeaf(GroovyTokenTypes.mCLOSABLE_BLOCK_OP, "->", next);
       getNode().addLeaf(GroovyTokenTypes.mNLS, "\n", next);
@@ -130,7 +135,7 @@ public class GrClosableBlockImpl extends GrBlockImpl implements GrClosableBlock
   }
 
   public boolean hasParametersSection() {
-    return findChildByType(GroovyElementTypes.mCLOSABLE_BLOCK_OP) != null;
+    return getArrow() != null;
   }
 
   public PsiType getType() {

plugins/groovy/test/org/jetbrains/plugins/groovy/intentions/GrIntentionTestCase.java

@@ -26,7 +26,7 @@ import java.util.List;
  * @author Maxim.Medvedev
  */
 public abstract class GrIntentionTestCase extends LightCodeInsightFixtureTestCase {
-  protected void doTest(String hint, boolean intentionExists) throws Exception {
+  protected void doTest(String hint, boolean intentionExists) {
     myFixture.configureByFile(getTestName(false) + ".groovy");
     final List<IntentionAction> list = myFixture.filterAvailableIntentions(hint);
     if (intentionExists) {

plugins/groovy/test/org/jetbrains/plugins/groovy/intentions/closure/eachToFor/EachToForIntentionTest.java

@@ -28,19 +28,23 @@ public class EachToForIntentionTest extends GrIntentionTestCase {
     return TestUtils.getTestDataPath() + "intentions/EachToFor/";
   }
 
-  public void testEachToFor() throws Throwable {
+  public void testEachToFor() {
     doTest("Replace with For-In", true);
   }
 
-  public void testEachToForWithFinal() throws Throwable {
+  public void testEachToForWithFinal() {
     doTest("Replace with For-In", true);
   }
 
-  public void testEachToForWithDefaultVariable() throws Throwable {
+  public void testEachToForWithDefaultVariable() {
     doTest("Replace with For-In", true);
   }
 
-  public void testEachForInWithNoQualifier () throws Throwable {
+  public void testEachForInWithNoQualifier() {
+    doTest("Replace with For-In", true);
+  }
+
+  public void testWithClosureInBody() {
     doTest("Replace with For-In", true);
   }
 }

plugins/groovy/testdata/intentions/EachToFor/WithClosureInBody.groovy

@@ -0,0 +1,10 @@
+[].ea<caret>ch {
+  if (it == 2) {
+    println 2
+  }
+  if (it == 3) {
+    println { String s ->
+      println s
+    }
+  }
+}

plugins/groovy/testdata/intentions/EachToFor/WithClosureInBody_after.groovy

@@ -0,0 +1,10 @@
+for (it in []) {
+  if (it == 2) {
+    println 2
+  }
+  if (it == 3) {
+    println { String s ->
+      println s
+    }
+  }
+}