Create a second IR between Pieces and output string. (#1517)

* Create a second IR between Pieces and output string.

Before this PR, CodeWriter eagerly built a complete string of formatted code using a StringBuffer. This is simple and works OK, but it's a little slow when separate formatting gets involved. Since a subtree of a Piece tree might be formatted separately, it means we often build a string using a StringBuffer only to then append that result to some surrounding StringBuffer, recursively.

This PR introduces a slightly more abstract representation for "formatted code but not quite a single string". It lets us compose separately formatted subtrees by just adding a single existing object to a list.

The performance gain isn't huge, but is measurable:

```
Benchmark (tall)                fastest   median  slowest  average  baseline
-----------------------------  --------  -------  -------  -------  --------
block                             0.063    0.065    0.114    0.069    101.0%
chain                             0.605    0.618    0.645    0.619    103.6%
collection                        0.159    0.163    0.175    0.164    101.4%
collection_large                  0.856    0.896    3.295    0.997    100.1%
conditional                       0.065    0.067    0.086    0.069     99.7%
curry                             0.547    0.565    0.596    0.565    106.8%
ffi                               0.144    0.150    0.170    0.151    110.0%
flutter_popup_menu_test           0.262    0.272    0.295    0.272    103.4%
flutter_scrollbar_test            0.123    0.125    0.142    0.128    103.7%
function_call                     1.287    1.339    1.474    1.342    101.1%
infix_large                       0.619    0.638    0.669    0.642    107.2%
infix_small                       0.159    0.171    0.719    0.228    105.5%
interpolation                     0.090    0.091    0.114    0.093    100.9%
large                             3.456    3.499    3.537    3.499    101.7%
top_level                         0.141    0.145    0.168    0.147    102.2%
```

And when run on a large repo:

```
Current formatter     8.536 ========================================
Optimized             8.226 ======================================
Old formatter         4.441 ====================
The current formatter is 47.97% slower than the old formatter.
The optimized is 3.77% faster than the current formatter.
The optimized is 46.01% slower than the old formatter.
The optimization gets the formatter 7.57% of the way to the old one.
```

The other motivation for this change is that I'm starting to work on being able to opt a region of code out of formatting and I suspect (but am not sure yet) that this representation will make it easier to slot the existing formatted code into the output as it's being lowered to a string.

* Fix typo.

Author: munificent

lib/src/back_end/code.dart

@@ -0,0 +1,174 @@
+// Copyright (c) 2024, 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.
+
+/// Base class for an object that represents fully formatted code.
+///
+/// We use this instead of immediately generating a string for the resulting
+/// formatted code because of separate formatting. Often, a subtree of the
+/// [Piece] tree can be solved and formatted separately. The resulting
+/// [Solution] may be used by multiple different surrounding solutions while
+/// the [Solver] works its magic looking for the best solution. When a
+/// separately formatted child solution is merged into its parent, we want that
+/// to be fast. Appending strings to a [StringBuffer] is fairly fast, but not
+/// as fast simply appending a single [GroupCode] to the parent solution's
+/// [GroupCode].
+sealed class Code {}
+
+/// A [Code] object which can be written to and contain other child [Code]
+/// objects.
+final class GroupCode extends Code {
+  /// The child [Code] objects contained in this group.
+  final List<Code> _children = [];
+
+  /// Appends [text] to this code.
+  void write(String text) {
+    _children.add(_TextCode(text));
+  }
+
+  /// Writes a newline and the subsequent indentation to this code.
+  ///
+  /// If [blank] is `true`, then a blank line is written. Otherwise, only a
+  /// single newline is written. The [indent] parameter is the number of spaces
+  /// of leading indentation on the next line after the newline.
+  void newline({required bool blank, required int indent}) {
+    _children.add(_NewlineCode(blank: blank, indent: indent));
+  }
+
+  /// Adds an entire existing code [group] as a child of this one.
+  void group(GroupCode group) {
+    _children.add(group);
+  }
+
+  /// Mark the selection start as occurring [offset] characters after the code
+  /// that has already been written.
+  void startSelection(int offset) {
+    _children.add(_MarkerCode(_Marker.start, offset));
+  }
+
+  /// Mark the selection end as occurring [offset] characters after the code
+  /// that has already been written.
+  void endSelection(int offset) {
+    _children.add(_MarkerCode(_Marker.end, offset));
+  }
+
+  /// Traverse the [Code] tree and build the final formatted string.
+  ///
+  /// Returns the formatted string and the selection markers if there are any.
+  ({String code, int? selectionStart, int? selectionEnd}) build() {
+    var buffer = StringBuffer();
+    int? selectionStart;
+    int? selectionEnd;
+
+    _build(buffer, (marker, offset) {
+      if (marker == _Marker.start) {
+        selectionStart = offset;
+      } else {
+        selectionEnd = offset;
+      }
+    });
+
+    return (
+      code: buffer.toString(),
+      selectionStart: selectionStart,
+      selectionEnd: selectionEnd
+    );
+  }
+
+  void _build(StringBuffer buffer,
+      void Function(_Marker marker, int offset) markSelection) {
+    for (var i = 0; i < _children.length; i++) {
+      var child = _children[i];
+      switch (child) {
+        case _NewlineCode():
+          // Don't write any leading newlines at the top of the buffer.
+          if (i > 0) {
+            buffer.writeln();
+            if (child._blank) buffer.writeln();
+          }
+
+          buffer.write(_indents[child._indent] ?? (' ' * child._indent));
+
+        case _TextCode():
+          buffer.write(child._text);
+
+        case GroupCode():
+          child._build(buffer, markSelection);
+
+        case _MarkerCode():
+          markSelection(child._marker, buffer.length + child._offset);
+      }
+    }
+  }
+}
+
+/// A [Code] object for a newline followed by any leading indentation.
+final class _NewlineCode extends Code {
+  final bool _blank;
+  final int _indent;
+
+  _NewlineCode({required bool blank, required int indent})
+      : _indent = indent,
+        _blank = blank;
+}
+
+/// A [Code] object for literal source text.
+final class _TextCode extends Code {
+  final String _text;
+
+  _TextCode(this._text);
+}
+
+/// Marks the location of the beginning or end of a selection as occurring
+/// [_offset] characters past the point where this marker object appears in the
+/// list of [Code] objects.
+final class _MarkerCode extends Code {
+  /// What kind of selection endpoint is being marked.
+  final _Marker _marker;
+
+  /// The number of characters past this object where the marker should appear
+  /// in the resulting code.
+  final int _offset;
+
+  _MarkerCode(this._marker, this._offset);
+}
+
+/// Which selection marker is pointed to by a [_MarkerCode].
+enum _Marker { start, end }
+
+/// Pre-calculated whitespace strings for various common levels of indentation.
+///
+/// Generating these ahead of time is faster than concatenating multiple spaces
+/// at runtime.
+const _indents = {
+  2: '  ',
+  4: '    ',
+  6: '      ',
+  8: '        ',
+  10: '          ',
+  12: '            ',
+  14: '              ',
+  16: '                ',
+  18: '                  ',
+  20: '                    ',
+  22: '                      ',
+  24: '                        ',
+  26: '                          ',
+  28: '                            ',
+  30: '                              ',
+  32: '                                ',
+  34: '                                  ',
+  36: '                                    ',
+  38: '                                      ',
+  40: '                                        ',
+  42: '                                          ',
+  44: '                                            ',
+  46: '                                              ',
+  48: '                                                ',
+  50: '                                                  ',
+  52: '                                                    ',
+  54: '                                                      ',
+  56: '                                                        ',
+  58: '                                                          ',
+  60: '                                                            ',
+};

lib/src/back_end/code_writer.dart

@@ -5,6 +5,7 @@ import 'dart:math';
 
 import '../piece/piece.dart';
 import '../profile.dart';
+import 'code.dart';
 import 'solution.dart';
 import 'solution_cache.dart';
 
@@ -28,8 +29,8 @@ class CodeWriter {
   /// The solution this [CodeWriter] is generating code for.
   final Solution _solution;
 
-  /// Buffer for the code being written.
-  final StringBuffer _buffer = StringBuffer();
+  /// The code being written.
+  final GroupCode _code = GroupCode();
 
   /// What whitespace should be written before the next non-whitespace text.
   ///
@@ -106,12 +107,12 @@ class CodeWriter {
     _pendingIndent = leadingIndent;
   }
 
-  /// Returns the final formatted text and the next pieces that can be expanded
+  /// Returns the final formatted code and the next pieces that can be expanded
   /// from the solution this [CodeWriter] is writing, if any.
-  (String, List<Piece>) finish() {
+  (GroupCode, List<Piece>) finish() {
     _finishLine();
 
-    return (_buffer.toString(), _expandPieces);
+    return (_code, _expandPieces);
   }
 
   /// Appends [text] to the output.
@@ -124,7 +125,7 @@ class CodeWriter {
   /// selections inside lexemes are correctly updated.
   void write(String text) {
     _flushWhitespace();
-    _buffer.write(text);
+    _code.write(text);
     _column += text.length;
 
     // If we haven't found an overflowing line yet, then this line might be one
@@ -253,20 +254,7 @@ class CodeWriter {
     _flushWhitespace();
 
     _solution.mergeSubtree(solution);
-
-    // If a selection marker was in the child piece, set it in this piece,
-    // relative to where the child's code is appended.
-    if (solution.selectionStart case var start?) {
-      _solution.startSelection(_buffer.length + start);
-    }
-
-    if (solution.selectionEnd case var end?) {
-      _solution.endSelection(_buffer.length + end);
-    }
-
-    Profile.begin('CodeWriter.format() write separate piece text');
-    _buffer.write(solution.text);
-    Profile.end('CodeWriter.format() write separate piece text');
+    _code.group(solution.code);
   }
 
   /// Format [piece] writing directly into this [CodeWriter].
@@ -334,13 +322,13 @@ class CodeWriter {
   /// Sets [selectionStart] to be [start] code units into the output.
   void startSelection(int start) {
     _flushWhitespace();
-    _solution.startSelection(_buffer.length + start);
+    _code.startSelection(start);
   }
 
   /// Sets [selectionEnd] to be [end] code units into the output.
   void endSelection(int end) {
     _flushWhitespace();
-    _solution.endSelection(_buffer.length + end);
+    _code.endSelection(end);
   }
 
   /// Write any pending whitespace.
@@ -355,18 +343,13 @@ class CodeWriter {
 
       case Whitespace.newline:
       case Whitespace.blankLine:
-        // Don't write any leading newlines at the top of the buffer.
-        if (_buffer.isNotEmpty) {
-          _finishLine();
-          _buffer.writeln();
-          if (_pendingWhitespace == Whitespace.blankLine) _buffer.writeln();
-        }
-
+        _finishLine();
         _column = _pendingIndent;
-        _buffer.write(' ' * _column);
+        _code.newline(
+            blank: _pendingWhitespace == Whitespace.blankLine, indent: _column);
 
       case Whitespace.space:
-        _buffer.write(' ');
+        _code.write(' ');
         _column++;
     }
 

lib/src/back_end/solution.dart

@@ -3,6 +3,7 @@
 // BSD-style license that can be found in the LICENSE file.
 import '../piece/piece.dart';
 import '../profile.dart';
+import 'code.dart';
 import 'code_writer.dart';
 import 'solution_cache.dart';
 
@@ -49,8 +50,8 @@ class Solution implements Comparable<Solution> {
   int _subtreeCost = 0;
 
   /// The formatted code.
-  String get text => _text;
-  late final String _text;
+  GroupCode get code => _code;
+  late final GroupCode _code;
 
   /// False if this Solution contains a newline where one is prohibited.
   ///
@@ -109,16 +110,6 @@ class Solution implements Comparable<Solution> {
   /// this one. It's either a dead end or a winner.
   late final List<Piece> _expandPieces;
 
-  /// The offset in [text] where the selection starts, or `null` if there is
-  /// no selection.
-  int? get selectionStart => _selectionStart;
-  int? _selectionStart;
-
-  /// The offset in [text] where the selection ends, or `null` if there is
-  /// no selection.
-  int? get selectionEnd => _selectionEnd;
-  int? _selectionEnd;
-
   /// Creates a new [Solution] with no pieces set to any state (which
   /// implicitly means they have state [State.unsplit] unless they're pinned to
   /// another state).
@@ -187,22 +178,6 @@ class Solution implements Comparable<Solution> {
     _subtreeCost += subtreeSolution.cost;
   }
 
-  /// Sets [selectionStart] to be [start] code units into the output.
-  ///
-  /// This should only be called by [CodeWriter].
-  void startSelection(int start) {
-    assert(_selectionStart == null);
-    _selectionStart = start;
-  }
-
-  /// Sets [selectionEnd] to be [end] code units into the output.
-  ///
-  /// This should only be called by [CodeWriter].
-  void endSelection(int end) {
-    assert(_selectionEnd == null);
-    _selectionEnd = end;
-  }
-
   /// Mark this solution as having a newline where none is permitted by [piece]
   /// and is thus not a valid solution.
   ///
@@ -319,9 +294,9 @@ class Solution implements Comparable<Solution> {
       SolutionCache cache, Piece root, int pageWidth, int leadingIndent) {
     var writer = CodeWriter(pageWidth, leadingIndent, cache, this);
     writer.format(root);
-    var (text, expandPieces) = writer.finish();
 
-    _text = text;
+    var (code, expandPieces) = writer.finish();
+    _code = code;
     _expandPieces = expandPieces;
   }
 

lib/src/back_end/solver.dart

@@ -103,7 +103,7 @@ class Solver {
 
       if (debug.traceSolver) {
         debug.log(debug.bold('Try #$attempts $solution'));
-        debug.log(solution.text);
+        debug.log(solution.code.build().code);
         debug.log('');
       }
 
@@ -134,7 +134,7 @@ class Solver {
     if (debug.traceSolver) {
       debug.unindent();
       debug.log(debug.bold('Solved $root to $best:'));
-      debug.log(best.text);
+      debug.log(solution.code.build().code);
       debug.log('');
     }