Revamp the way tokens and comments are built into pieces. (#1336)

Revamp the way tokens and comments are built into pieces.

I recently ran into bugs where a line comment after some AST node will cause the node to split incorrectly. A simple example is:

```
var x = 1 + 2; // comment
```

Currently, the formatter splits that to:

```
var x =
    1 +
        2; // comment
```

It does that because the Piece tree it creates doesn't line up with the AST node boundaries. In particular, the current design appends tokens and comments to a preceding piece. So in this example, the piece tree looks like:

```
Var(
  `var`
  Assign(
    `x =`
    Infix(
      `1 +`
      `2; // comment`
    )
  )
)
```

Note how the `;` and line comment are attached as part of the RHS of the `+`. That's why the formatter thinks the line comment's newline is inside the + expression and forces it to split.

We could fix this specific bug by making ExpressionStatement treat the `;` as a separate piece, but I suspect that we will be playing whack-a-mole if we keep the current design. Instead, this unfortunately giant PR revamps the piece API. It has a couple of intermingled changes:

### Split pieces at all AST boundaries

Whenever a `visit___()` returns, an implicit split is inserted so that no single `TextPiece` contains tokens from a parent and child AST node. This directly fixes the above bug and all similar bugs in that category.

Note that while we split the tokens into separate pieces, that doesn't mean they may split in the "line splitting" sense. The TextPieces go into AdjacentPiece objects that don't insert actual splits between the pieces. This means this change shouldn't significantly impact the performance of line splitting.

It's just about ensuring that the nesting structure of the piece tree mirrors the nesting structure of the AST. That way, when a newline in a child piece node invalidates an outer piece, that invalidation respecst the original syntax.

### Revamp the API for creating pieces

The previous API had a DSL-like "push" API where the pieces created by PieceWriter were stored internally and exposed by a fairly confusing `give()`/`take()`/`split()` API. That was necessary because any given `visit___()` method might not be *able* to return a Piece for its node if that node just concatenated its tokens into some surrounding piece.

With the previous change where every AST node corresponds to a piece, we have that option. So this PR also makes that change. Every `visit___()` method is now required to return a piece. Likewise, all of the `create___()` methods in PieceFactory return the pieces they create. This avoids the need for a weird `take()` API.

### Add an AdjacentBuilder and buildPiece() API

Getting rid of the implicit storage and dataflow for pieces is good for being able to easily reason about how the piece tree gets created out of child pieces. But it can come at the cost of making code that creates pieces very verbose with lots of local variables and `List<Piece>` objects to store the intermediate pieces being built.

To make that nicer, I wrote an AdjacentBuilder class with an imperative API for building an AdjacentPiece out of a series of tokens, nodes, and spaces. This API closely mirrors the original DSL-like API. Except now you know exactly what object the nodes and tokens are pushing their pieces into.

To make that even nicer, I added a `buildPiece()` method that takes a callback, invokes it with a new AdjacentBuilder, and return the built result. This gets most code for building pieces fairly close to the original push-based API but with hopefully clearer more explicit dataflow.

I'm really sorry for the giant size of this PR. If you want, I can try to break it into a series of smaller commits (but likely still one PR), but doing so is pretty challenging given how intertwined these changes are. It's hard to change the return type of the visit methods without also getting rid of the implicit dataflow and at that point, almost all the changes are there.

Also, I added more tests to cover the cases around comments that were broken.

Author: munificent

lib/src/back_end/code_writer.dart

@@ -1,6 +1,8 @@
 // Copyright (c) 2023, the Dart project authors.  Please see the AUTHORS file
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
+import 'dart:math';
+
 import '../piece/piece.dart';
 import 'solution.dart';
 
@@ -34,11 +36,11 @@ class CodeWriter {
   /// it as pending. This ensures that we don't write trailing whitespace,
   /// avoids writing spaces at the beginning of lines, and allows collapsing
   /// multiple redundant newlines.
-  _Whitespace _pendingWhitespace = _Whitespace.none;
+  Whitespace _pendingWhitespace = Whitespace.none;
 
   /// The number of spaces of indentation that should be begin the next line
-  /// when [_pendingWhitespace] is [_Whitespace.newline] or
-  /// [_Whitespace.blankLine].
+  /// when [_pendingWhitespace] is [Whitespace.newline] or
+  /// [Whitespace.blankLine].
   int _pendingIndent = 0;
 
   /// The cost of the currently chosen line splits.
@@ -194,10 +196,7 @@ class CodeWriter {
 
   /// Writes a single space to the output.
   void space() {
-    // If a newline is already pending, then ignore the space.
-    if (_pendingWhitespace == _Whitespace.none) {
-      _pendingWhitespace = _Whitespace.space;
-    }
+    whitespace(Whitespace.space);
   }
 
   /// Inserts a line split in the output.
@@ -209,16 +208,16 @@ class CodeWriter {
   void newline({bool blank = false, int? indent}) {
     if (indent != null) setIndent(indent);
 
-    handleNewline();
+    whitespace(blank ? Whitespace.blankLine : Whitespace.newline);
+  }
 
-    // Collapse redundant newlines.
-    if (blank) {
-      _pendingWhitespace = _Whitespace.blankLine;
-    } else if (_pendingWhitespace != _Whitespace.blankLine) {
-      _pendingWhitespace = _Whitespace.newline;
+  void whitespace(Whitespace whitespace) {
+    if (whitespace case Whitespace.newline || Whitespace.blankLine) {
+      handleNewline();
+      _pendingIndent = _options.indent;
     }
 
-    _pendingIndent = _options.indent;
+    _pendingWhitespace = _pendingWhitespace.collapse(whitespace);
   }
 
   /// Sets whether newlines are allowed to occur from this point on for the
@@ -286,24 +285,24 @@ class CodeWriter {
   /// count of the written text, including whitespace.
   void _flushWhitespace() {
     switch (_pendingWhitespace) {
-      case _Whitespace.none:
+      case Whitespace.none:
         break; // Nothing to do.
 
-      case _Whitespace.newline:
-      case _Whitespace.blankLine:
+      case Whitespace.newline:
+      case Whitespace.blankLine:
         _finishLine();
         _buffer.writeln();
-        if (_pendingWhitespace == _Whitespace.blankLine) _buffer.writeln();
+        if (_pendingWhitespace == Whitespace.blankLine) _buffer.writeln();
 
         _column = _pendingIndent;
         _buffer.write(' ' * _column);
 
-      case _Whitespace.space:
+      case Whitespace.space:
         _buffer.write(' ');
         _column++;
     }
 
-    _pendingWhitespace = _Whitespace.none;
+    _pendingWhitespace = Whitespace.none;
   }
 
   void _finishLine() {
@@ -328,18 +327,28 @@ class CodeWriter {
 }
 
 /// Different kinds of pending whitespace that have been requested.
-enum _Whitespace {
+///
+/// Note that the order of values in the enum is significant: later ones have
+/// more whitespace than previous ones.
+enum Whitespace {
   /// No pending whitespace.
   none,
 
+  /// A single space.
+  space,
+
   /// A single newline.
   newline,
 
   /// Two newlines.
-  blankLine,
+  blankLine;
 
-  /// A single space.
-  space
+  /// Combines two pending whitespaces and returns the result.
+  ///
+  /// When two whitespaces overlap, they aren't both written: we don't want
+  /// two spaces or a newline followed by a space. Instead, the two whitespaces
+  /// are collapsed such that the largest one wins.
+  Whitespace collapse(Whitespace other) => values[max(index, other.index)];
 }
 
 /// The mutable state local to a single piece being formatted.

lib/src/front_end/adjacent_builder.dart

@@ -0,0 +1,109 @@
+// Copyright (c) 2023, the Dart project authors.  Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+import 'package:analyzer/dart/ast/ast.dart';
+import 'package:analyzer/dart/ast/token.dart';
+
+import '../piece/adjacent.dart';
+import '../piece/piece.dart';
+import 'piece_factory.dart';
+
+/// Incrementally builds an [AdjacentPiece].
+class AdjacentBuilder {
+  final PieceFactory _visitor;
+
+  /// The series of adjacent pieces.
+  final List<Piece> _pieces = [];
+
+  AdjacentBuilder(this._visitor);
+
+  /// Yields a new piece containing all of the pieces added to or created by
+  /// this builder. The caller must ensure it doesn't build an empty piece.
+  ///
+  /// Also clears the builder's list of pieces so that this builder can be
+  /// reused to build more pieces.
+  Piece build() {
+    assert(_pieces.isNotEmpty);
+
+    var result = _flattenPieces();
+    _pieces.clear();
+
+    return result;
+  }
+
+  /// Adds [piece] to this builder.
+  void add(Piece piece) {
+    _pieces.add(piece);
+  }
+
+  /// Emit [token], along with any comments and formatted whitespace that comes
+  /// before it.
+  ///
+  /// If [lexeme] is given, uses that for the token's lexeme instead of its own.
+  ///
+  /// Does nothing if [token] is `null`. If [spaceBefore] is `true`, writes a
+  /// space before the token, likewise with [spaceAfter].
+  void token(Token? token,
+      {bool spaceBefore = false, bool spaceAfter = false, String? lexeme}) {
+    if (token == null) return;
+
+    if (spaceBefore) space();
+    add(_visitor.pieces.tokenPiece(token, lexeme: lexeme));
+    if (spaceAfter) space();
+  }
+
+  /// Writes any comments that appear before [token], which will be discarded.
+  ///
+  /// Used to ensure comments before a discarded token are preserved.
+  void commentsBefore(Token? token) {
+    if (token == null) return;
+
+    var piece = _visitor.pieces.writeCommentsBefore(token);
+    if (piece != null) add(piece);
+  }
+
+  /// Writes an optional modifier that precedes other code.
+  void modifier(Token? keyword) {
+    token(keyword, spaceAfter: true);
+  }
+
+  /// Visits [node] if not `null` and adds the resulting [Piece] to this
+  /// builder.
+  void visit(AstNode? node,
+      {bool spaceBefore = false,
+      bool commaAfter = false,
+      bool spaceAfter = false}) {
+    if (node == null) return;
+
+    if (spaceBefore) space();
+    add(_visitor.nodePiece(node, commaAfter: commaAfter));
+    if (spaceAfter) space();
+  }
+
+  /// Appends a space before the previous piece and the next one.
+  void space() {
+    _pieces.add(SpacePiece());
+  }
+
+  /// Removes redundant [AdjacentPiece] wrappers from [_pieces].
+  Piece _flattenPieces() {
+    List<Piece> flattened = [];
+
+    void traverse(List<Piece> pieces) {
+      for (var piece in pieces) {
+        if (piece is AdjacentPiece) {
+          traverse(piece.pieces);
+        } else {
+          flattened.add(piece);
+        }
+      }
+    }
+
+    traverse(_pieces);
+
+    // If there's only one piece, don't wrap it in a pointless AdjacentPiece.
+    if (flattened.length == 1) return flattened[0];
+
+    return AdjacentPiece(flattened);
+  }
+}