Extract Function: Allow to select fragments which start with comment

idea/src/org/jetbrains/jet/plugin/codeInsight/CodeInsightUtils.java

@@ -97,18 +97,11 @@ public class CodeInsightUtils {
         }
         if (endOffset != element2.getTextRange().getEndOffset()) return PsiElement.EMPTY_ARRAY;
 
-        PsiElement[] children = parent.getChildren();
         ArrayList<PsiElement> array = new ArrayList<PsiElement>();
-        boolean flag = false;
+        PsiElement stopElement = element2.getNextSibling();
-        for (PsiElement child : children) {
+        for (PsiElement currentElement = element1; currentElement != stopElement; currentElement = currentElement.getNextSibling()) {
-            if (child.equals(element1)) {
+            if (!(currentElement instanceof PsiWhiteSpace)) {
-                flag = true;
+                array.add(currentElement);
-            }
-            if (flag && !(child instanceof PsiWhiteSpace)) {
-                array.add(child);
-            }
-            if (child.equals(element2)) {
-                break;
             }
         }
 

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/ExtractionData.kt

@@ -106,7 +106,7 @@ class ExtractionData(
     }
 
     fun getBrokenReferencesInfo(body: JetBlockExpression): List<ResolvedReferenceInfo> {
-        val startOffset = body.getStatements().first!!.getTextRange()!!.getStartOffset()
+        val startOffset = body.getBlockContentOffset()
 
         val referencesInfo = ArrayList<ResolvedReferenceInfo>()
         for ((ref, context) in JetFileReferencesResolver.resolve(body)) {
@@ -139,3 +139,10 @@ private fun compareDescriptors(d1: DeclarationDescriptor?, d2: DeclarationDescri
             (d1 != null && d2 != null &&
                     DescriptorRenderer.FQ_NAMES_IN_TYPES.render(d1) == DescriptorRenderer.FQ_NAMES_IN_TYPES.render(d2))
 }
+
+// Hack:
+// we can't get first element offset through getStatement()/getChildren() since they skip comments and whitespaces
+// So we take offset of the left brace instead and increase it by 2 (which is length of "{\n" separating block start and its first element)
+private fun JetBlockExpression.getBlockContentOffset(): Int {
+    return getLBrace()!!.getTextRange()!!.getStartOffset() + 2
+}

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/extractFunctionUtils.kt

@@ -681,10 +681,7 @@ fun ExtractionDescriptor.generateFunction(
         val exprReplacementMap = HashMap<JetElement, (JetElement) -> JetElement>()
         val originalOffsetByExpr = HashMap<JetElement, Int>()
 
-        val range = body.getStatements().first?.getTextRange()
+        val bodyOffset = body.getBlockContentOffset()
-        if (range == null) return
-
-        val bodyOffset = range.getStartOffset()
         val file = body.getContainingFile()!!
 
         for ((offsetInBody, resolveResult) in extractionData.refOffsetToDeclaration) {

idea/testData/refactoring/extractFunction/basic/fragmentWithComment.kt

@@ -0,0 +1,7 @@
+// NEXT_SIBLING:
+fun foo(a: Int): Int {
+    <selection>// test
+    println(a)
+    if (a > 0) return a</selection>
+    return -a
+}

idea/testData/refactoring/extractFunction/basic/fragmentWithComment.kt.after

@@ -0,0 +1,12 @@
+// NEXT_SIBLING:
+fun b(a: Int): Boolean {
+    // test
+    println(a)
+    if (a > 0) return true
+    return false
+}
+
+fun foo(a: Int): Int {
+    if (b(a)) return a
+    return -a
+}

idea/testData/refactoring/extractFunction/basic/fragmentWithMultilineComment.kt

@@ -0,0 +1,9 @@
+// NEXT_SIBLING:
+fun foo(a: Int): Int {
+    <selection>/*
+    test
+    */
+    println(a)
+    if (a > 0) return a</selection>
+    return -a
+}

idea/testData/refactoring/extractFunction/basic/fragmentWithMultilineComment.kt.after

@@ -0,0 +1,14 @@
+// NEXT_SIBLING:
+fun b(a: Int): Boolean {
+    /*
+        test
+        */
+    println(a)
+    if (a > 0) return true
+    return false
+}
+
+fun foo(a: Int): Int {
+    if (b(a)) return a
+    return -a
+}

idea/tests/org/jetbrains/jet/plugin/refactoring/introduce/introduceVariable/JetExtractionTestGenerated.java

@@ -198,6 +198,16 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                 JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.TestsPackage", new File("idea/testData/refactoring/extractFunction/basic"), Pattern.compile("^(.+)\\.kt$"), true);
             }
             
+            @TestMetadata("fragmentWithComment.kt")
+            public void testFragmentWithComment() throws Exception {
+                doExtractFunctionTest("idea/testData/refactoring/extractFunction/basic/fragmentWithComment.kt");
+            }
+            
+            @TestMetadata("fragmentWithMultilineComment.kt")
+            public void testFragmentWithMultilineComment() throws Exception {
+                doExtractFunctionTest("idea/testData/refactoring/extractFunction/basic/fragmentWithMultilineComment.kt");
+            }
+            
             @TestMetadata("localClassFunctionRef.kt")
             public void testLocalClassFunctionRef() throws Exception {
                 doExtractFunctionTest("idea/testData/refactoring/extractFunction/basic/localClassFunctionRef.kt");