Get selections running through the new back end. (#1300)

Get selections running through the new back end.

This mostly ports similar code directly over from the old back end,
including a couple of quirks in its behavior around whitespace. I
considered simplifying that, but I think it's better to have fewer
changes during this already huge migration.

Because the new back end adds and removes trailing commas, it treats
all commas as whitespace. When a selection is in whitespace, it will be
snapped to the nearest non-whitespace code. These two interact in that
a selection right around a comma may snap over it because the comma is
considered whitespace.

I thought about making it smarter but I didn't want to rathole too long.
Instead, I left a comment in CodeWriter.write() with some context and
we can revisit it when the overall rewrite is further along. We may end
up just deleting that comment later if we decide the behavior here is
good enough (I think it is), but this way to remember to think about it.

Author: munificent

lib/src/back_end/code_writer.dart

@@ -59,16 +59,33 @@ class CodeWriter {
   /// The options for the current innermost piece being formatted.
   _PieceOptions get _options => _pieceOptions.last;
 
+  /// The offset in the formatted code where the selection starts.
+  ///
+  /// This is `null` until the piece containing the selection start is reached
+  /// at which point it gets set. It remains `null` if there is no selection.
+  int? _selectionStart;
+
+  /// The offset in the formatted code where the selection ends.
+  ///
+  /// This is `null` until the piece containing the selection end is reached
+  /// at which point it gets set. It remains `null` if there is no selection.
+  int? _selectionEnd;
+
   CodeWriter(this._pageWidth, this._pieceStates);
 
   /// Returns the finished code produced by formatting the tree of pieces and
   /// the final score.
-  (String, Score) finish() {
+  Solution finish() {
     _finishLine();
-    return (
-      _buffer.toString(),
-      Score(isValid: !_containsInvalidNewline, overflow: _overflow, cost: _cost)
-    );
+    return Solution(
+        _pieceStates,
+        _buffer.toString(),
+        Score(
+            isValid: !_containsInvalidNewline,
+            overflow: _overflow,
+            cost: _cost),
+        _selectionStart,
+        _selectionEnd);
   }
 
   /// Notes that a newline has been written.
@@ -92,6 +109,14 @@ class CodeWriter {
   /// If [text] contains any internal newlines, the caller is responsible for
   /// also calling [handleNewline()].
   void write(String text) {
+    // TODO(tall): Calling this directly from pieces outside of TextPiece may
+    // not handle selections as gracefully as we could. A selection marker may
+    // get pushed past the text written here. Currently, this is only called
+    // directly for commas in list-like things, and `;` in for loops. In
+    // general, it's better for all text written to the output to live inside
+    // TextPieces because that will preserve selection markers. Consider doing
+    // something smarter for commas in lists and semicolons in for loops.
+
     _buffer.write(text);
     _column += text.length;
   }
@@ -204,6 +229,18 @@ class CodeWriter {
     if (piece != null) format(piece);
   }
 
+  /// Sets [selectionStart] to be [start] code units into the output.
+  void startSelection(int start) {
+    assert(_selectionStart == null);
+    _selectionStart = _buffer.length + start;
+  }
+
+  /// Sets [selectionEnd] to be [end] code units into the output.
+  void endSelection(int end) {
+    assert(_selectionEnd == null);
+    _selectionEnd = _buffer.length + end;
+  }
+
   void _finishLine() {
     // If the completed line is too long, track the overflow.
     if (_column >= _pageWidth) {

lib/src/back_end/solution.dart

@@ -72,14 +72,26 @@ class Solution implements Comparable<Solution> {
   /// The score resulting from the selected piece states.
   final Score score;
 
-  factory Solution(Piece root, int pageWidth, PieceStateSet state) {
+  /// The offset in [text] where the selection starts, or `null` if there is
+  /// no selection.
+  final int? selectionStart;
+
+  /// The offset in [text] where the selection ends, or `null` if there is
+  /// no selection.
+  final int? selectionEnd;
+
+  factory Solution.initial(Piece root, int pageWidth, List<Piece> pieces) {
+    return Solution._(root, pageWidth, PieceStateSet(pieces));
+  }
+
+  factory Solution._(Piece root, int pageWidth, PieceStateSet state) {
     var writer = CodeWriter(pageWidth, state);
     writer.format(root);
-    var (text, score) = writer.finish();
-    return Solution._(state, text, score);
+    return writer.finish();
   }
 
-  Solution._(this._state, this.text, this.score);
+  Solution(this._state, this.text, this.score, this.selectionStart,
+      this.selectionEnd);
 
   /// When called on a [Solution] with some unselected piece states, chooses a
   /// piece and yields further solutions for each state that piece can have.
@@ -89,7 +101,7 @@ class Solution implements Comparable<Solution> {
 
     return [
       for (var state in piece.states)
-        Solution(root, pageWidth, _state.cloneWith(piece, state))
+        Solution._(root, pageWidth, _state.cloneWith(piece, state))
     ];
   }
 

lib/src/back_end/solver.dart

@@ -32,7 +32,7 @@ class Solver {
 
   /// Finds the best set of line splits for [piece] and returns the resulting
   /// formatted code.
-  String format(Piece piece) {
+  Solution format(Piece piece) {
     // Collect all of the pieces with states that can be selected.
     var unsolvedPieces = <Piece>[];
 
@@ -44,14 +44,13 @@ class Solver {
 
     traverse(piece);
 
-    var solution = _solve(piece, unsolvedPieces);
-    return solution.text;
+    return _solve(piece, unsolvedPieces);
   }
 
   /// Finds the best solution for the piece tree starting at [root] with
   /// selectable [pieces].
   Solution _solve(Piece root, List<Piece> pieces) {
-    var solution = Solution(root, _pageWidth, PieceStateSet(pieces));
+    var solution = Solution.initial(root, _pageWidth, pieces);
     _queue.add(solution);
 
     // The lowest cost solution found so far that does overflow.

lib/src/front_end/ast_node_visitor.dart

@@ -913,11 +913,9 @@ class AstNodeVisitor extends ThrowingAstVisitor<void>
   @override
   void visitScriptTag(ScriptTag node) {
     // The lexeme includes the trailing newline. Strip it off since the
-    // formatter ensures it gets a newline after it. Since the script tag must
-    // come at the top of the file, we don't have to worry about preceding
-    // comments or whitespace.
-    // TODO(new-ir): Update selection if inside the script tag.
-    pieces.write(node.scriptTag.lexeme.trim());
+    // formatter ensures it gets a newline after it.
+    pieces.writeText(node.scriptTag.lexeme.trim(),
+        offset: node.scriptTag.offset);
   }
 
   @override

lib/src/front_end/comment_writer.dart

@@ -141,10 +141,8 @@ mixin CommentWriter {
         type = CommentType.block;
       }
 
-      var sourceComment = SourceComment(text, type, flushLeft: flushLeft);
-
-      // TODO(tall): If this comment contains either of the selection endpoints,
-      // mark them in the comment.
+      var sourceComment = SourceComment(text, type,
+          offset: comment.offset, flushLeft: flushLeft);
 
       comments._add(linesBefore, sourceComment);
 
@@ -194,7 +192,14 @@ class SourceComment {
   /// re-indented.
   final bool flushLeft;
 
-  SourceComment(this.text, this.type, {required this.flushLeft});
+  /// The number of code points in the original source code preceding the start
+  /// of this comment.
+  ///
+  /// Used to track selection markers within the comment.
+  final int offset;
+
+  SourceComment(this.text, this.type,
+      {required this.flushLeft, required this.offset});
 
   /// Whether this comment contains a mandatory newline, either because it's a
   /// comment that should be on its own line or a lexeme with a newline inside

lib/src/front_end/piece_factory.dart

@@ -470,16 +470,10 @@ mixin PieceFactory implements CommentWriter {
     writeCommentsBefore(token);
 
     if (before != null) before();
-    writeLexeme(token.lexeme);
+    pieces.writeToken(token);
     if (after != null) after();
   }
 
-  /// Writes the raw [lexeme] to the current text piece.
-  void writeLexeme(String lexeme) {
-    // TODO(tall): Preserve selection.
-    pieces.write(lexeme);
-  }
-
   /// Writes a comma after [node], if there is one.
   void commaAfter(AstNode node, {bool trailing = false}) {
     var nextToken = node.endToken.next!;
@@ -489,7 +483,7 @@ mixin PieceFactory implements CommentWriter {
       // If there isn't a comma there, it must be a place where a trailing
       // comma can appear, so synthesize it. During formatting, we will decide
       // whether to include it.
-      writeLexeme(',');
+      pieces.writeText(',');
     }
   }
 }