Implement range focus with special case for lists.

Author: toinehartman

rascal-lsp/src/main/java/org/rascalmpl/vscode/lsp/parametric/ParametricTextDocumentService.java

@@ -770,12 +770,7 @@ public CompletableFuture<List<? extends TextEdit>> rangeFormatting(DocumentRange
                 var start = Locations.toRascalPosition(uri, range.getStart(), columns);
                 var end = Locations.toRascalPosition(uri, range.getEnd(), columns);
                 // compute the focus list at the end of the range
-                var focus = TreeSearch.computeFocusList(tree, end.getLine(), end.getCharacter())
-                    .stream()
-                    .map(ITree.class::cast)
-                    // check for containment of the start of the range
-                    .filter(t -> Ranges.containsPosition(Locations.toRange(TreeAdapter.getLocation(t), columns), start))
-                    .collect(VF.listWriter());
+                var focus = TreeSearch.computeFocusList(tree, start.getLine(), start.getCharacter(), end.getLine(), end.getCharacter());
 
                 var opts = getFormattingOptions(params.getOptions());
                 return contribs.formatting(focus, opts).get();

rascal-lsp/src/main/java/org/rascalmpl/vscode/lsp/util/locations/impl/TreeSearch.java

@@ -26,6 +26,8 @@
  */
 package org.rascalmpl.vscode.lsp.util.locations.impl;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.rascalmpl.values.IRascalValueFactory;
 import org.rascalmpl.values.parsetrees.ITree;
 import org.rascalmpl.values.parsetrees.TreeAdapter;
@@ -40,6 +42,9 @@
  */
 public class TreeSearch {
 
+    private static final Logger logger = LogManager.getLogger(TreeSearch.class);
+    private static final IRascalValueFactory VF = IRascalValueFactory.getInstance();
+
     private TreeSearch() {}
 
     /**
@@ -85,6 +90,17 @@ else if (line == loc.getEndLine()) {
         return true;
     }
 
+    private static boolean rightOf(ISourceLocation loc, int line, int column) {
+        if (!loc.hasLineColumn()) {
+            return false;
+        }
+
+        if (line > loc.getEndLine()) {
+            return true;
+        }
+        return line == loc.getEndLine() && column > loc.getEndColumn();
+    }
+
     /**
      * Produces a list of trees that are "in focus" at given line and column offset (UTF-24).
      *
@@ -99,7 +115,7 @@ else if (line == loc.getEndLine()) {
      * @return list of tree that are around the given line/column position, ordered from child to parent.
      */
     public static IList computeFocusList(ITree tree, int line, int column) {
-        var lw = IRascalValueFactory.getInstance().listWriter();
+        var lw = VF.listWriter();
         computeFocusList(lw, tree, line, column);
         return lw.done();
     }
@@ -157,4 +173,61 @@ private static boolean computeFocusList(IListWriter focus, ITree tree, int line,
         // cycles and characters do not have locations
         return false;
     }
+
+    public static IList computeFocusList(ITree tree, int startLine, int startColumn, int endLine, int endColumn) {
+        // Compute the focus for both the start end end positions.
+        // These foci give us information about the structure of the selection.
+        final var startList = computeFocusList(tree, startLine, startColumn);
+        final var endList = computeFocusList(tree, endLine, endColumn);
+
+        final var commonSuffix = startList.intersect(endList);
+        if (commonSuffix.equals(startList) || commonSuffix.equals(endList)) {
+            // We do not have enough information to extend the focus
+            return commonSuffix;
+        }
+        // The range spans multiple subtrees. The easy way out is not to focus farther down than
+        // their smallest common subtree (i.e. `commonSuffix`) - let's see if we can do any better.
+        if (TreeAdapter.isList((ITree) commonSuffix.get(0))) {
+            return computeListRangeFocus(commonSuffix, startLine, startColumn, endLine, endColumn);
+        }
+
+        return commonSuffix;
+    }
+
+    private static IList computeListRangeFocus(final IList commonSuffix, int startLine, int startColumn, int endLine, int endColumn) {
+        final var parent = (ITree) commonSuffix.get(0);
+        logger.trace("Computing focus list for {} at range [{}:{}, {}:{}]", TreeAdapter.getType(parent), startLine, startColumn, endLine, endColumn);
+        final var elements = TreeAdapter.getListASTArgs(parent);
+        final int nElements = elements.length();
+
+        logger.trace("Smallest common tree is a {} with {} elements", TreeAdapter.getType(parent), nElements);
+        if (inside(TreeAdapter.getLocation((ITree) elements.get(0)), startLine, startColumn) &&
+            inside(TreeAdapter.getLocation((ITree) elements.get(nElements - 1)), endLine, endColumn)) {
+            // The whole list is selected
+            return commonSuffix;
+        }
+
+        // Find the elements in the list that are (partially) selected.
+        final var selected = elements.stream()
+            .map(ITree.class::cast)
+            .dropWhile(t -> !inside(TreeAdapter.getLocation(t), startLine, startColumn))
+            .takeWhile(t -> rightOf(TreeAdapter.getLocation(t), endLine, endColumn))
+            .collect(VF.listWriter());
+        final int nSelected = selected.length();
+
+        logger.trace("Range covers {} (of {}) elements in the parent list", nSelected, nElements);
+        final var firstSelected = TreeAdapter.getLocation((ITree) selected.get(0));
+        final var lastSelected = TreeAdapter.getLocation((ITree) selected.get(nSelected - 1));
+
+        final int totalLength = lastSelected.getOffset() - firstSelected.getOffset() + lastSelected.getLength();
+        final var selectionLoc = VF.sourceLocation(firstSelected, firstSelected.getOffset(), totalLength,
+            firstSelected.getBeginLine(), lastSelected.getEndLine(), firstSelected.getBeginColumn(), lastSelected.getEndColumn());
+        final var artificialParent = TreeAdapter.setLocation(VF.appl(TreeAdapter.getProduction(parent), selected), selectionLoc);
+
+        // Build new focus list
+        var lw = VF.listWriter();
+        lw.append(artificialParent);
+        lw.appendAll(commonSuffix);
+        return lw.done();
+    }
 }