Fix movement bug when expression is followed/preceded by comment at the functional literal block boundary

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

@@ -2,14 +2,19 @@ package org.jetbrains.jet.plugin.codeInsight.upDownMover;
 
 import com.intellij.codeInsight.editorActions.moveUpDown.LineMover;
 import com.intellij.codeInsight.editorActions.moveUpDown.LineRange;
+import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.util.Pair;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiComment;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.lang.psi.JetBlockExpression;
 import org.jetbrains.jet.lang.psi.JetFile;
+import org.jetbrains.jet.lang.psi.JetFunctionLiteral;
 
 public abstract class AbstractJetUpDownMover extends LineMover {
     protected AbstractJetUpDownMover() {
@@ -31,6 +36,40 @@ public abstract class AbstractJetUpDownMover extends LineMover {
         }
 
         PsiElement parent = PsiTreeUtil.findCommonParent(firstElement, lastElement);
+
+        int topExtension = 0;
+        int bottomExtension = 0;
+
+        if (parent instanceof JetFunctionLiteral) {
+            JetBlockExpression block = ((JetFunctionLiteral) parent).getBodyExpression();
+            if (block != null) {
+                PsiElement comment = null;
+
+                boolean extendDown = false;
+                if (checkCommentAtBlockBound(firstElement, lastElement, block)) {
+                    comment = lastElement;
+                    extendDown = true;
+                    lastElement = block.getLastChild();
+                } else if (checkCommentAtBlockBound(lastElement, firstElement, block)) {
+                    comment = firstElement;
+                    firstElement = block.getFirstChild();
+                }
+
+                if (comment != null) {
+                    int extension = getElementLineCount(comment, editor);
+                    if (extendDown) {
+                        bottomExtension = extension;
+                    }
+                    else {
+                        topExtension = extension;
+                    }
+                }
+
+
+                parent = PsiTreeUtil.findCommonParent(firstElement, lastElement);
+            }
+        }
+
         if (parent == null) return null;
 
         Pair<PsiElement, PsiElement> combinedRange = getElementRange(parent, firstElement, lastElement);
@@ -49,7 +88,7 @@ public abstract class AbstractJetUpDownMover extends LineMover {
 
         LineRange parentLineRange = getElementSourceLineRange(parent, editor, oldRange);
 
-        LineRange sourceRange = new LineRange(lineRange1.startLine, lineRange2.endLine);
+        LineRange sourceRange = new LineRange(lineRange1.startLine - topExtension, lineRange2.endLine + bottomExtension);
         if (parentLineRange != null && sourceRange.contains(parentLineRange)) {
             sourceRange.firstElement = sourceRange.lastElement = parent;
         } else {
@@ -60,6 +99,20 @@ public abstract class AbstractJetUpDownMover extends LineMover {
         return sourceRange;
     }
 
+    private static int getElementLineCount(PsiElement element, Editor editor) {
+        Document doc = editor.getDocument();
+        TextRange spaceRange = element.getTextRange();
+
+        int startLine = doc.getLineNumber(spaceRange.getStartOffset());
+        int endLine = doc.getLineNumber(spaceRange.getEndOffset());
+
+        return endLine - startLine;
+    }
+
+    private static boolean checkCommentAtBlockBound(PsiElement blockElement, PsiElement comment, JetBlockExpression block) {
+        return PsiTreeUtil.isAncestor(block, blockElement, true) && comment instanceof PsiComment;
+    }
+
     @Nullable
     protected static PsiElement getSiblingOfType(@NotNull PsiElement element, boolean down, @NotNull Class<? extends PsiElement> type) {
         return down ? PsiTreeUtil.getNextSiblingOfType(element, type) : PsiTreeUtil.getPrevSiblingOfType(element, type);

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

@@ -0,0 +1,8 @@
+// MOVE: up
+fun foo() {
+    bar {
+        /**/ <caret>val foo = 1
+        println("foo=")
+        println(foo) /**/
+    }
+}

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

@@ -0,0 +1,8 @@
+// MOVE: up
+fun foo() {
+    /**/ <caret>val foo = 1
+    bar {
+        println("foo=")
+        println(foo) /**/
+    }
+}

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

@@ -0,0 +1,8 @@
+// MOVE: down
+fun foo() {
+    bar {
+        /**/ val foo = 1
+        println("foo=")
+        <caret>println(foo) /**/
+    }
+}

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

@@ -0,0 +1,8 @@
+// MOVE: down
+fun foo() {
+    bar {
+        /**/ val foo = 1
+        println("foo=")
+    }
+    <caret>println(foo) /**/
+}

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

@@ -777,6 +777,16 @@ 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("closureBlockBoundary1.kt")
+        public void testClosureBlockBoundary1() throws Exception {
+            doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/closureBlockBoundary1.kt");
+        }
+        
+        @TestMetadata("closureBlockBoundary2.kt")
+        public void testClosureBlockBoundary2() throws Exception {
+            doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/closureBlockBoundary2.kt");
+        }
+        
         @TestMetadata("closureInDoWhile1.kt")
         public void testClosureInDoWhile1() throws Exception {
             doTestExpression("idea/testData/codeInsight/moveUpDown/expressions/closureInDoWhile1.kt");