Fix "Move statement" for the cases when nested closure is used (KT-3735)

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiUtil.java

@@ -779,22 +779,22 @@ public class JetPsiUtil {
     }
 
     @Nullable
-    public static <T extends JetElement> T getOutermostJetElement(
+    public static JetElement getOutermostDescendantElement(
             @Nullable PsiElement root,
             boolean first,
-            @NotNull final Class<? extends T>... elementTypes
+            final @NotNull Predicate<JetElement> predicate
     ) {
         if (!(root instanceof JetElement)) return null;
 
-        final List<T> results = Lists.newArrayList();
+        final List<JetElement> results = Lists.newArrayList();
 
         ((JetElement) root).accept(
                 new JetVisitorVoid() {
                     @Override
                     public void visitJetElement(JetElement element) {
-                        if (PsiTreeUtil.instanceOf(element, elementTypes)) {
+                        if (predicate.apply(element)) {
                             //noinspection unchecked
-                            results.add((T) element);
+                            results.add(element);
                         }
                         else {
                             element.acceptChildren(this);

idea/src/org/jetbrains/jet/plugin/codeInsight/upDownMover/JetExpressionMover.java

@@ -1,5 +1,6 @@
 package org.jetbrains.jet.plugin.codeInsight.upDownMover;
 
+import com.google.common.base.Predicate;
 import com.intellij.codeInsight.editorActions.moveUpDown.LineRange;
 import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.Editor;
@@ -13,10 +14,12 @@ import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lexer.JetTokens;
 
 public class JetExpressionMover extends AbstractJetUpDownMover {
+
     public JetExpressionMover() {
     }
 
-    private final static Class[] MOVABLE_ELEMENT_CLASSES = {JetExpression.class, JetWhenEntry.class, JetValueArgument.class, PsiComment.class};
+    private final static Class[] MOVABLE_ELEMENT_CLASSES =
+            {JetExpression.class, JetWhenEntry.class, JetValueArgument.class, PsiComment.class};
 
     private final static Class[] BLOCKLIKE_ELEMENT_CLASSES =
             {JetBlockExpression.class, JetWhenExpression.class, JetClassBody.class, JetFile.class};
@@ -24,6 +27,21 @@ public class JetExpressionMover extends AbstractJetUpDownMover {
     private final static Class[] FUNCTIONLIKE_ELEMENT_CLASSES =
             {JetFunction.class, JetPropertyAccessor.class, JetClassInitializer.class};
 
+    private static final Predicate<JetElement> CHECK_BLOCK_LIKE_ELEMENT = new Predicate<JetElement>() {
+        @Override
+        public boolean apply(@Nullable JetElement input) {
+            return (input instanceof JetBlockExpression || input instanceof JetClassBody)
+                   && !PsiTreeUtil.instanceOf(input.getParent(), FUNCTIONLIKE_ELEMENT_CLASSES);
+        }
+    };
+
+    private static final Predicate<JetElement> CHECK_BLOCK = new Predicate<JetElement>() {
+        @Override
+        public boolean apply(@Nullable JetElement input) {
+            return input instanceof JetBlockExpression && !PsiTreeUtil.instanceOf(input.getParent(), FUNCTIONLIKE_ELEMENT_CLASSES);
+        }
+    };
+
     @Nullable
     private static PsiElement getStandaloneClosingBrace(@NotNull PsiFile file, @NotNull Editor editor) {
         LineRange range = getLineRangeFromSelection(editor);
@@ -137,7 +155,7 @@ public class JetExpressionMover extends AbstractJetUpDownMover {
             PsiElement parent = current.getParent();
             if (parent instanceof JetClassBody ||
                 parent instanceof JetClassInitializer ||
-                parent instanceof JetFunction ||
+                parent instanceof JetNamedFunction ||
                 parent instanceof JetProperty) {
                 return null;
             }
@@ -147,7 +165,7 @@ public class JetExpressionMover extends AbstractJetUpDownMover {
             PsiElement sibling = down ? current.getNextSibling() : current.getPrevSibling();
             if (sibling != null) {
                 //noinspection unchecked
-                JetBlockExpression block = JetPsiUtil.getOutermostJetElement(sibling, down, JetBlockExpression.class);
+                JetBlockExpression block = (JetBlockExpression) JetPsiUtil.getOutermostDescendantElement(sibling, down, CHECK_BLOCK);
                 if (block != null) return block;
 
                 current = sibling;
@@ -195,21 +213,13 @@ public class JetExpressionMover extends AbstractJetUpDownMover {
         else {
             // moving into code block
             //noinspection unchecked
-            JetElement blockLikeElement = JetPsiUtil.getOutermostJetElement(sibling, down, JetBlockExpression.class, JetWhenExpression.class, JetClassBody.class);
+            JetElement blockLikeElement = JetPsiUtil.getOutermostDescendantElement(sibling, down, CHECK_BLOCK_LIKE_ELEMENT);
-            if (blockLikeElement != null &&
+            if (blockLikeElement != null) {
-                !(PsiTreeUtil.instanceOf(blockLikeElement.getParent(), FUNCTIONLIKE_ELEMENT_CLASSES))) {
+                if (down) {
-                if (blockLikeElement instanceof JetWhenExpression) {
+                    end = JetPsiUtil.findChildByType(blockLikeElement, JetTokens.LBRACE);
-                    //noinspection unchecked
-                    blockLikeElement = JetPsiUtil.getOutermostJetElement(blockLikeElement, down, JetBlockExpression.class);
                 }
-
+                else {
-                if (blockLikeElement != null) {
+                    start = JetPsiUtil.findChildByType(blockLikeElement, JetTokens.RBRACE);
-                    if (down) {
-                        end = JetPsiUtil.findChildByType(blockLikeElement, JetTokens.LBRACE);
-                    }
-                    else {
-                        start = JetPsiUtil.findChildByType(blockLikeElement, JetTokens.RBRACE);
-                    }
                 }
             }
         }
@@ -228,7 +238,12 @@ public class JetExpressionMover extends AbstractJetUpDownMover {
     }
 
     @Nullable
-    private LineRange getValueParamOrArgTargetRange(@NotNull Editor editor, @NotNull PsiElement elementToCheck, @NotNull PsiElement sibling, boolean down) {
+    private LineRange getValueParamOrArgTargetRange(
+            @NotNull Editor editor,
+            @NotNull PsiElement elementToCheck,
+            @NotNull PsiElement sibling,
+            boolean down
+    ) {
         PsiElement next = sibling;
 
         if (next.getNode().getElementType() == JetTokens.COMMA) {
@@ -236,8 +251,8 @@ public class JetExpressionMover extends AbstractJetUpDownMover {
         }
 
         LineRange range = (next instanceof JetParameter || next instanceof JetValueArgument)
-               ? new LineRange(next, next, editor.getDocument())
+                          ? new LineRange(next, next, editor.getDocument())
-               : null;
+                          : null;
 
         if (range != null) {
             parametersOrArgsToMove = new Pair<PsiElement, PsiElement>(elementToCheck, next);
@@ -374,8 +389,10 @@ public class JetExpressionMover extends AbstractJetUpDownMover {
                 info.toMove2 = null;
                 return true;
             }
-            case MOVABLE: return true;
+            case MOVABLE:
-            default: break;
+                return true;
+            default:
+                break;
         }
 
         LineRange oldRange = info.toMove;
@@ -394,7 +411,8 @@ public class JetExpressionMover extends AbstractJetUpDownMover {
             return true;
         }
 
-        if ((firstElement instanceof JetParameter || firstElement instanceof JetValueArgument) && PsiTreeUtil.isAncestor(lastElement, firstElement, false)) {
+        if ((firstElement instanceof JetParameter || firstElement instanceof JetValueArgument) &&
+            PsiTreeUtil.isAncestor(lastElement, firstElement, false)) {
             lastElement = firstElement;
         }
 

idea/testData/codeInsight/moveUpDown/expressions/closureInDoWhile1.kt

@@ -0,0 +1,8 @@
+// MOVE: down
+fun foo(i: Int) {
+    do {
+        println(i)
+        <caret>run {
+        }
+    } while (i in run { 1..2 })
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInDoWhile1.kt.after

@@ -0,0 +1,8 @@
+// MOVE: down
+fun foo(i: Int) {
+    do {
+        println(i)
+    } while (i in run { 1..2 })
+    <caret>run {
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInDoWhile2.kt

@@ -0,0 +1,8 @@
+// MOVE: up
+fun foo(i: Int) {
+    do {
+        println(i)
+    } while (i in run { 1..2 })
+    <caret>run {
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInDoWhile2.kt.after

@@ -0,0 +1,8 @@
+// MOVE: up
+fun foo(i: Int) {
+    do {
+        println(i)
+        <caret>run {
+        }
+    } while (i in run { 1..2 })
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInFor1.kt

@@ -0,0 +1,7 @@
+// MOVE: down
+fun foo() {
+    <caret>run {
+    }
+    for (i in run { 1..2 }) {
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInFor1.kt.after

@@ -0,0 +1,7 @@
+// MOVE: down
+fun foo() {
+    for (i in run { 1..2 }) {
+        <caret>run {
+        }
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInFor2.kt

@@ -0,0 +1,7 @@
+// MOVE: up
+fun foo() {
+    for (i in run { 1..2 }) {
+        <caret>run {
+        }
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInFor2.kt.after

@@ -0,0 +1,7 @@
+// MOVE: up
+fun foo() {
+    <caret>run {
+    }
+    for (i in run { 1..2 }) {
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInIf1.kt

@@ -0,0 +1,7 @@
+// MOVE: down
+fun foo(i: Int) {
+    <caret>run {
+    }
+    if (i in run { 1..2 }) {
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInIf1.kt.after

@@ -0,0 +1,7 @@
+// MOVE: down
+fun foo(i: Int) {
+    if (i in run { 1..2 }) {
+        <caret>run {
+        }
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInIf2.kt

@@ -0,0 +1,7 @@
+// MOVE: up
+fun foo(i: Int) {
+    if (i in run { 1..2 }) {
+        <caret>run {
+        }
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInIf2.kt.after

@@ -0,0 +1,7 @@
+// MOVE: up
+fun foo(i: Int) {
+    <caret>run {
+    }
+    if (i in run { 1..2 }) {
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInWhile1.kt

@@ -0,0 +1,8 @@
+// MOVE: down
+fun foo(i: Int) {
+    <caret>run {
+    }
+    while (i in run { 1..2 }) {
+        println(i)
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInWhile1.kt.after

@@ -0,0 +1,8 @@
+// MOVE: down
+fun foo(i: Int) {
+    while (i in run { 1..2 }) {
+        <caret>run {
+        }
+        println(i)
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInWhile2.kt

@@ -0,0 +1,8 @@
+// MOVE: up
+fun foo(i: Int) {
+    while (i in run { 1..2 }) {
+        <caret>run {
+        }
+        println(i)
+    }
+}

idea/testData/codeInsight/moveUpDown/expressions/closureInWhile2.kt.after

@@ -0,0 +1,8 @@
+// MOVE: up
+fun foo(i: Int) {
+    <caret>run {
+    }
+    while (i in run { 1..2 }) {
+        println(i)
+    }
+}

idea/tests/org/jetbrains/jet/plugin/codeInsight/moveUpDown/CodeMoverTestGenerated.java

@@ -777,6 +777,46 @@ public class CodeMoverTestGenerated extends AbstractCodeMoverTest {
             JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.GenerateTests", new File("idea/testData/codeInsight/moveUpDown/expressions"), Pattern.compile("^(.+)\\.kt$"), true);
         }
         
+        @TestMetadata("closureInDoWhile1.kt")
+        public void testClosureInDoWhile1() throws Exception {
+            doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/closureInDoWhile1.kt");
+        }
+        
+        @TestMetadata("closureInDoWhile2.kt")
+        public void testClosureInDoWhile2() throws Exception {
+            doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/closureInDoWhile2.kt");
+        }
+        
+        @TestMetadata("closureInFor1.kt")
+        public void testClosureInFor1() throws Exception {
+            doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/closureInFor1.kt");
+        }
+        
+        @TestMetadata("closureInFor2.kt")
+        public void testClosureInFor2() throws Exception {
+            doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/closureInFor2.kt");
+        }
+        
+        @TestMetadata("closureInIf1.kt")
+        public void testClosureInIf1() throws Exception {
+            doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/closureInIf1.kt");
+        }
+        
+        @TestMetadata("closureInIf2.kt")
+        public void testClosureInIf2() throws Exception {
+            doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/closureInIf2.kt");
+        }
+        
+        @TestMetadata("closureInWhile1.kt")
+        public void testClosureInWhile1() throws Exception {
+            doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/closureInWhile1.kt");
+        }
+        
+        @TestMetadata("closureInWhile2.kt")
+        public void testClosureInWhile2() throws Exception {
+            doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/closureInWhile2.kt");
+        }
+        
         @TestMetadata("declaration1.kt")
         public void testDeclaration1() throws Exception {
             doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/declaration1.kt");