PY-11357 Make automatic surrounding with custom folding comments more robust

* Searching of first parental PSI elements before/after new line
correctly stops at the start/end of file. As result it's possible now
to surround element at the first/last line in the file.
Also corresponding traversal does not try to climb above PSI file.
* Even single character can be selected in first place.
* In languages that uses indentation to delimit blocks (like Python)
several consequent sibling statements can be surrounded even if the last
of them is at the end of its parent.

platform/lang-api/src/com/intellij/lang/folding/CustomFoldingSurroundDescriptor.java

@@ -26,10 +26,7 @@ import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.util.text.StringUtil;
-import com.intellij.psi.PsiDocumentManager;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiWhiteSpace;
+import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
 import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
@@ -62,7 +59,7 @@ public class CustomFoldingSurroundDescriptor implements SurroundDescriptor {
   @NotNull
   @Override
   public PsiElement[] getElementsToSurround(PsiFile file, int startOffset, int endOffset) {
-    if (startOffset >= endOffset - 1) return PsiElement.EMPTY_ARRAY;
+    if (startOffset >= endOffset) return PsiElement.EMPTY_ARRAY;
     Commenter commenter = LanguageCommenters.INSTANCE.forLanguage(file.getLanguage());
     if (commenter == null || commenter.getLineCommentPrefix() == null) return PsiElement.EMPTY_ARRAY;
     PsiElement startElement = file.findElementAt(startOffset);
@@ -70,18 +67,11 @@ public class CustomFoldingSurroundDescriptor implements SurroundDescriptor {
     PsiElement endElement = file.findElementAt(endOffset - 1);
     if (endElement instanceof PsiWhiteSpace) endElement = endElement.getPrevSibling();
     if (startElement != null && endElement != null) {
-      if (startElement.getTextRange().getStartOffset() > endElement.getTextRange().getStartOffset()) return PsiElement.EMPTY_ARRAY;
       startElement = findClosestParentAfterLineBreak(startElement);
       if (startElement != null) {
         endElement = findClosestParentBeforeLineBreak(endElement);
         if (endElement != null) {
-          startElement = adjustStartElementIfEndAbsorbed(startElement, endElement);
-          endElement = adjustEndElementIfStartAbsorbed(startElement, endElement);
-          final PsiElement commonParent = startElement.getParent();
-          if (endElement.getParent() == commonParent) {
-            if (startElement == endElement) return new PsiElement[]{startElement};
-            return new PsiElement[]{startElement, endElement};
-          }
+          return adjustRange(startElement, endElement);
         }
       }
     }
@@ -89,46 +79,116 @@ public class CustomFoldingSurroundDescriptor implements SurroundDescriptor {
   }
 
   @NotNull
-  private static PsiElement adjustEndElementIfStartAbsorbed(@NotNull PsiElement start, @NotNull PsiElement end) {
-    if (PsiTreeUtil.isAncestor(end, start, false) && start.getTextRange().getEndOffset() == end.getTextRange().getEndOffset()) {
-      return start;
+  private static PsiElement[] adjustRange(@NotNull PsiElement start, @NotNull PsiElement end) {
+    PsiElement newStart = lowerStartElementIfNeeded(start, end);
+    PsiElement newEnd = lowerEndElementIfNeeded(start, end);
+    if (newStart == null || newEnd == null) {
+      return PsiElement.EMPTY_ARRAY;
+    }
+    final PsiElement commonParent = findCommonAncestorForWholeRange(newStart, newEnd);
+    if (commonParent != null) {
+      return new PsiElement[] {commonParent};
+    }
+    // If either start or end element is the first/last leaf element in its parent, use the parent itself instead
+    // to prevent selection of clearly illegal ranges like the following:
+    // [
+    //   <selection>1
+    // ]</selection>
+    // E.g. in case shown, because of that adjustment, closing bracket and number literal won't have the same parent
+    // and next test will fail.
+    if (newStart.getParent().getFirstChild() == newStart && newStart.getFirstChild() == null) {
+      newStart = newStart.getParent();
+    }
+    if (newEnd.getParent().getLastChild() == newEnd && newEnd.getFirstChild() == null) {
+      newEnd = newEnd.getParent();
+    }
+    if (newStart.getParent() == newEnd.getParent()) {
+      return new PsiElement[] {newStart, newEnd};
+    }
+    return PsiElement.EMPTY_ARRAY;
+  }
+
+  @Nullable
+  private static PsiElement lowerEndElementIfNeeded(@NotNull PsiElement start, @NotNull PsiElement end) {
+    if (PsiTreeUtil.isAncestor(end, start, true)) {
+      PsiElement lastChild = end.getLastChild();
+      while (lastChild != null && lastChild.getParent() != start.getParent()) {
+        lastChild = lastChild.getLastChild();
+      }
+      return lastChild;
     }
     return end;
   }
 
-  @NotNull
-  private static PsiElement adjustStartElementIfEndAbsorbed(@NotNull PsiElement start, @NotNull PsiElement end) {
-    if (PsiTreeUtil.isAncestor(start, end, false) && start.getTextRange().getStartOffset() == end.getTextRange().getStartOffset()) {
-      return end;
+  @Nullable
+  private static PsiElement lowerStartElementIfNeeded(@NotNull PsiElement start, @NotNull PsiElement end) {
+    if (PsiTreeUtil.isAncestor(start, end, true)) {
+      PsiElement firstChild = start.getFirstChild();
+      while (firstChild != null && firstChild.getParent() != end.getParent()) {
+        firstChild = firstChild.getFirstChild();
+      }
+      return firstChild;
     }
     return start;
   }
 
   @Nullable
-  private static PsiElement findClosestParentAfterLineBreak(PsiElement element) {
-    PsiElement parent = element;
-    while (parent != null) {
-      PsiElement prev = parent.getPrevSibling();
-      while (prev != null && prev.getTextLength() <= 0) {
-        prev = prev.getPrevSibling();
-      }
-      if (isWhiteSpaceWithLineFeed(prev)) return parent;
-      parent = parent.getParent();
+  private static PsiElement findCommonAncestorForWholeRange(@NotNull PsiElement start, @NotNull PsiElement end) {
+    final PsiElement parent = PsiTreeUtil.findCommonParent(start, end);
+    if (parent == null) {
+      return null;
+    }
+    final TextRange parentRange = parent.getTextRange();
+    if (parentRange.getStartOffset() == start.getTextRange().getStartOffset() &&
+        parentRange.getEndOffset() == end.getTextRange().getEndOffset()) {
+      return parent;
     }
     return null;
   }
 
   @Nullable
-  private static PsiElement findClosestParentBeforeLineBreak(PsiElement element) {
+  private static PsiElement findClosestParentAfterLineBreak(PsiElement element) {
     PsiElement parent = element;
-    while (parent != null) {
-      PsiElement next = parent.getNextSibling();
-      if (isWhiteSpaceWithLineFeed(next)) return parent;
+    while (parent != null && !(parent instanceof PsiFileSystemItem)) {
+      PsiElement prev = parent.getPrevSibling();
+      while (prev != null && prev.getTextLength() <= 0) {
+        prev = prev.getPrevSibling();
+      }
+      if (firstElementInFile(parent)) {
+        return parent.getContainingFile();
+      }
+      else if (isWhiteSpaceWithLineFeed(prev)) {
+        return parent;
+      }
       parent = parent.getParent();
     }
     return null;
   }
 
+  private static boolean firstElementInFile(@NotNull PsiElement element) {
+    return element.getTextOffset() == 0;
+  }
+
+  @Nullable
+  private static PsiElement findClosestParentBeforeLineBreak(PsiElement element) {
+    PsiElement parent = element;
+    while (parent != null && !(parent instanceof PsiFileSystemItem)) {
+      final PsiElement next = parent.getNextSibling();
+      if (lastElementInFile(parent)) {
+        return parent.getContainingFile();
+      }
+      else if (isWhiteSpaceWithLineFeed(next)) {
+        return parent;
+      }
+      parent = parent.getParent();
+    }
+    return null;
+  }
+
+  private static boolean lastElementInFile(@NotNull PsiElement element) {
+    return element.getTextRange().getEndOffset() == element.getContainingFile().getTextRange().getEndOffset();
+  }
+
   private static boolean isWhiteSpaceWithLineFeed(@Nullable PsiElement element) {
     if (element == null) {
       return false;

plugins/groovy/test/org/jetbrains/plugins/groovy/lang/surroundWith/SurrounderOrderTest.groovy

@@ -39,7 +39,9 @@ class SurrounderOrderTest extends LightCodeInsightFixtureTestCase {
                         "(expr)",
                         "!(expr)",
                         "((Type) expr)",
-                        "with () {...}"
+                        "with () {...}",
+                        "<editor-fold...> Comments",
+                        "region...endregion Comments"
   }
 
   public void testStatementWithSemicolon() throws Exception {
@@ -50,7 +52,9 @@ class SurrounderOrderTest extends LightCodeInsightFixtureTestCase {
                         "{}",
                         "for", "try / catch", "try / finally", "try / catch / finally",
                         "shouldFail () {...}",
-                        "with () {...}"
+                        "with () {...}",
+                        "<editor-fold...> Comments",
+                        "region...endregion Comments"
   }
 
   public void testStatementsWithComments() throws Exception {
@@ -64,7 +68,9 @@ println c /*also important */
                         "{}",
                         "for", "try / catch", "try / finally", "try / catch / finally",
                         "shouldFail () {...}",
-                        "with () {...}"
+                        "with () {...}",
+                        "<editor-fold...> Comments",
+                        "region...endregion Comments"
   }
 
   public void testInnerExpressionSurrounders() {

python/testData/surround/CustomFoldingRegionFirstMethod_after.py

@@ -0,0 +1,17 @@
+class ThisIsATest():
+    # <editor-fold desc="Description">
+    def __init__(self):
+        self.test = 1
+    # </editor-fold>
+
+    def another_one(self):
+        print "Hello, world!"
+
+    def another_two(self):
+        print "Hello, world!"
+
+    def another_three(self):
+        print "Hello, world!"
+
+    def another_four(self):
+        print "Hello, world!"

python/testData/surround/CustomFoldingRegionIllegalSelection.py

@@ -0,0 +1,4 @@
+[
+  1,
+  <selection>2
+]</selection>

python/testData/surround/CustomFoldingRegionIllegalSelection_after.py

@@ -0,0 +1,4 @@
+[
+  1,
+  2
+]

python/testData/surround/CustomFoldingRegionLastMethod.py

@@ -12,4 +12,4 @@ class ThisIsATest():
         print "Hello, world!"
 
     <selection>def another_four(self):
-        print "Hello, world!"</selection>
+        print "Hello, world!"</selection>

python/testData/surround/CustomFoldingRegionLastMethod_after.py

@@ -0,0 +1,17 @@
+class ThisIsATest():
+    def __init__(self):
+        self.test = 1
+
+    def another_one(self):
+        print "Hello, world!"
+
+    def another_two(self):
+        print "Hello, world!"
+
+    def another_three(self):
+        print "Hello, world!"
+
+    # <editor-fold desc="Description">
+    def another_four(self):
+        print "Hello, world!"
+        # </editor-fold>

python/testData/surround/CustomFoldingRegionSeveralMethods.py

@@ -0,0 +1,9 @@
+class C:
+    def m1(self):
+        pass
+
+    <selection>def m2(self):
+        pass
+
+    def m3(self):
+        pass</selection>

python/testData/surround/CustomFoldingRegionSeveralMethods_after.py

@@ -0,0 +1,11 @@
+class C:
+    def m1(self):
+        pass
+
+    # <editor-fold desc="Description">
+    def m2(self):
+        pass
+
+    def m3(self):
+        pass
+        # </editor-fold>

python/testData/surround/CustomFoldingRegionSingleCharacter.py

@@ -0,0 +1 @@
+x = 'foo' <selection>+</selection> 'bar'

python/testData/surround/CustomFoldingRegionSingleCharacter_after.py

@@ -0,0 +1,3 @@
+# <editor-fold desc="Description">
+x = 'foo' + 'bar'
+# </editor-fold>

python/testData/surround/CustomFoldingRegionSingleStatementInFile.py

@@ -0,0 +1,2 @@
+<selection>print('foo')
+</selection>

python/testData/surround/CustomFoldingRegionSingleStatementInFile_after.py

@@ -0,0 +1,3 @@
+# <editor-fold desc="Description">
+print('foo')
+# </editor-fold>

python/testSrc/com/jetbrains/python/PySurroundWithTest.java

@@ -19,11 +19,9 @@ import com.intellij.codeInsight.generation.surroundWith.SurroundWithHandler;
 import com.intellij.lang.folding.CustomFoldingSurroundDescriptor;
 import com.intellij.lang.surroundWith.Surrounder;
 import com.intellij.openapi.command.WriteCommandAction;
-import com.intellij.openapi.editor.SelectionModel;
-import com.intellij.psi.PsiElement;
+import com.intellij.openapi.util.Condition;
+import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.fixtures.PyTestCase;
-import com.jetbrains.python.psi.PyElement;
-import com.jetbrains.python.psi.PyFunction;
 import com.jetbrains.python.refactoring.surround.surrounders.statements.PyWithIfSurrounder;
 import com.jetbrains.python.refactoring.surround.surrounders.statements.PyWithTryExceptSurrounder;
 import com.jetbrains.python.refactoring.surround.surrounders.statements.PyWithWhileSurrounder;
@@ -45,32 +43,45 @@ public class PySurroundWithTest extends PyTestCase {
   }
 
   // PY-11357
-  public void testSurroundFirstMethodWithCustomFoldingRegion() {
-    checkCustomFoldingRegionRange(PyFunction.class);
+  public void testCustomFoldingRegionFirstMethod() throws Exception {
+    doTestSurroundWithCustomFoldingRegion();
   }
 
   // PY-11357
-  public void testSurroundLastMethodWithCustomFoldingRegion() {
-    checkCustomFoldingRegionRange(PyFunction.class);
+  public void testCustomFoldingRegionLastMethod() throws Exception {
+    doTestSurroundWithCustomFoldingRegion();
   }
 
   // PY-14261
-  public void testSurroundWithCustomFoldingRegion() throws Exception {
-    doTest(CustomFoldingSurroundDescriptor.SURROUNDERS[0]);
+  public void testCustomFoldingRegionPreservesIndentation() throws Exception {
+    doTestSurroundWithCustomFoldingRegion();
   }
 
+  public void testCustomFoldingRegionSingleCharacter() throws Exception {
+    doTestSurroundWithCustomFoldingRegion();
+  }
 
-  private PsiElement[] checkCustomFoldingRegionRange(Class<? extends PyElement>... elementTypes) {
-    myFixture.configureByFile("/surround/" + getTestName(false) + ".py");
-    final SelectionModel selection = myFixture.getEditor().getSelectionModel();
-    final PsiElement[] range = CustomFoldingSurroundDescriptor.INSTANCE.getElementsToSurround(myFixture.getFile(),
-                                                                                              selection.getSelectionStart(),
-                                                                                              selection.getSelectionEnd());
-    assertEquals(elementTypes.length, range.length);
-    for (int i = 0; i < elementTypes.length; i++) {
-      assertInstanceOf(range[i], elementTypes[i]);
-    }
-    return range;
+  public void testCustomFoldingRegionSingleStatementInFile() throws Exception {
+    doTestSurroundWithCustomFoldingRegion();
+  }
+
+  public void testCustomFoldingRegionIllegalSelection() throws Exception {
+    doTestSurroundWithCustomFoldingRegion();
+  }
+
+  public void testCustomFoldingRegionSeveralMethods() throws Exception {
+    doTestSurroundWithCustomFoldingRegion();
+  }
+
+  private void doTestSurroundWithCustomFoldingRegion() throws Exception {
+    final Surrounder surrounder = ContainerUtil.find(CustomFoldingSurroundDescriptor.SURROUNDERS, new Condition<Surrounder>() {
+      @Override
+      public boolean value(Surrounder surrounder) {
+        return surrounder.getTemplateDescription().contains("<editor-fold");
+      }
+    });
+    assertNotNull(surrounder);
+    doTest(surrounder);
   }
 
   private void doTest(final Surrounder surrounder) throws Exception {