Prune solutions derived from invalid ones. (#1372)

Prune solutions derived from invalid ones.

A solution is invalid if a child piece writes a newline at a point in
time where a parent disallows newlines.

When that happens, the parent's may be explicitly bound to a state by
the solution, or it may still be in an unbound state (and thus using
the default). If unbound, then even though this solution is invalid, it
may lead to later valid ones by binding that parent piece to a different
state.

But once a solution binds a piece to a state, all solutions derived
from that state keep the same binding. So if we hit an invalid solution
and the piece whose newline constraint was already violated, we can
discard that solution and all possible solutions derived from it.

This skips an entire potentially huge subtree from the solution space
being explored.

In an example giant deeply nested function call from the benchmark,
before this optimization, the solver tries 329,057 solutions before
finding a winning one. With this optimization, it only takes 5,405.

Author: munificent

lib/src/back_end/code_writer.dart

@@ -17,15 +17,11 @@ import 'solution.dart';
 /// an instance of this class. It has methods that the piece can call to add
 /// output text to the resulting code, recursively format child pieces, insert
 /// whitespace, etc.
-///
-/// This class also accumulates the score (the relative desireability of a set
-/// of formatting choices) that the resulting code has by tracking things like
-/// how many characters of code overflow the page width.
 class CodeWriter {
   final int _pageWidth;
 
-  /// The state values for the pieces being written.
-  final PieceStateSet _pieceStates;
+  /// The solution this [CodeWriter] is generating code for.
+  final Solution _solution;
 
   /// Buffer for the code being written.
   final StringBuffer _buffer = StringBuffer();
@@ -43,20 +39,9 @@ class CodeWriter {
   /// [Whitespace.blankLine].
   int _pendingIndent = 0;
 
-  /// The cost of the currently chosen line splits.
-  int _cost = 0;
-
-  /// The total number of characters of code that have overflowed the page
-  /// width so far.
-  int _overflow = 0;
-
   /// The number of characters in the line currently being written.
   int _column = 0;
 
-  /// Whether this solution has encountered a mandatory newline (like from a
-  /// line comment or a statement terminator) where no newline is permitted.
-  bool _hasInvalidNewline = false;
-
   /// The stack of state for each [Piece] being formatted.
   ///
   /// For each piece being formatted from a call to [format()], we keep track of
@@ -67,7 +52,7 @@ class CodeWriter {
   ///
   /// This is used to increase the cumulative nesting as we recurse into pieces
   /// and then unwind that as child pieces are completed.
-  final List<_PieceOptions> _pieceOptions = [_PieceOptions(0, true)];
+  final List<_PieceOptions> _options = [];
 
   /// Whether we have already found the first line where whose piece should be
   /// used to expand further solutions.
@@ -101,31 +86,14 @@ class CodeWriter {
   /// solution if the line ends up overflowing.
   final List<Piece> _currentUnsolvedPieces = [];
 
-  /// 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);
+  CodeWriter(this._pageWidth, this._solution);
 
-  /// Returns the finished code produced by formatting the tree of pieces and
-  /// the final score.
-  Solution finish() {
+  /// Returns the final formatted text and the next piece that can be expanded
+  /// from the solution this [CodeWriter] is writing, if any.
+  (String, Piece?) finish() {
     _finishLine();
 
-    return Solution(_pieceStates, _buffer.toString(), _selectionStart,
-        _selectionEnd, _nextPieceToExpand,
-        isValid: !_hasInvalidNewline, overflow: _overflow, cost: _cost);
+    return (_buffer.toString(), _nextPieceToExpand);
   }
 
   /// Appends [text] to the output.
@@ -158,10 +126,17 @@ class CodeWriter {
   /// piece to [indent], relative to the indentation of the surrounding piece.
   ///
   /// Replaces any previous indentation set by this piece.
-  ///
   // TODO(tall): Add another API that adds/subtracts existing indentation.
   void setIndent(int indent) {
-    _options.indent = _pieceOptions[_pieceOptions.length - 2].indent + indent;
+    var parentIndent = 0;
+
+    // If there is a surrounding Piece, then set the indent relative to that
+    // piece's current indentation.
+    if (_options.length > 1) {
+      parentIndent = _options[_options.length - 2].indent;
+    }
+
+    _options.last.indent = parentIndent + indent;
   }
 
   /// Inserts a newline if [condition] is true.
@@ -215,7 +190,7 @@ class CodeWriter {
   void whitespace(Whitespace whitespace, {bool flushLeft = false}) {
     if (whitespace case Whitespace.newline || Whitespace.blankLine) {
       _handleNewline();
-      _pendingIndent = flushLeft ? 0 : _options.indent;
+      _pendingIndent = flushLeft ? 0 : _options.last.indent;
     }
 
     _pendingWhitespace = _pendingWhitespace.collapse(whitespace);
@@ -224,31 +199,29 @@ class CodeWriter {
   /// Sets whether newlines are allowed to occur from this point on for the
   /// current piece.
   void setAllowNewlines(bool allowed) {
-    _options.allowNewlines = allowed;
+    _options.last.allowNewlines = allowed;
   }
 
   /// Format [piece] and insert the result into the code being written and
   /// returned by [finish()].
   void format(Piece piece) {
-    _pieceOptions.add(_PieceOptions(_options.indent, _options.allowNewlines));
+    _options.add(_PieceOptions(piece, _options.lastOrNull?.indent ?? 0,
+        _options.lastOrNull?.allowNewlines ?? true));
 
-    var isUnsolved = !_pieceStates.isBound(piece) && piece.states.length > 1;
+    var isUnsolved = !_solution.isBound(piece) && piece.states.length > 1;
     if (isUnsolved) _currentUnsolvedPieces.add(piece);
 
-    var state = _pieceStates.pieceState(piece);
-    _cost += piece.stateCost(state);
-
     // TODO(perf): Memoize this. Might want to create a nested PieceWriter
     // instead of passing in `this` so we can better control what state needs
     // to be used as the key in the memoization table.
-    piece.format(this, state);
+    piece.format(this, _solution.pieceState(piece));
 
     if (isUnsolved) _currentUnsolvedPieces.removeLast();
 
-    var childOptions = _pieceOptions.removeLast();
+    var childOptions = _options.removeLast();
 
     // If the child [piece] contains a newline then this one transitively does.
-    if (childOptions.hasNewline) _handleNewline();
+    if (childOptions.hasNewline && _options.isNotEmpty) _handleNewline();
   }
 
   /// Format [piece] if not null.
@@ -258,30 +231,26 @@ class CodeWriter {
 
   /// Sets [selectionStart] to be [start] code units into the output.
   void startSelection(int start) {
-    assert(_selectionStart == null);
-
     _flushWhitespace();
-    _selectionStart = _buffer.length + start;
+    _solution.startSelection(_buffer.length + start);
   }
 
   /// Sets [selectionEnd] to be [end] code units into the output.
   void endSelection(int end) {
-    assert(_selectionEnd == null);
-
     _flushWhitespace();
-    _selectionEnd = _buffer.length + end;
+    _solution.endSelection(_buffer.length + end);
   }
 
   /// Notes that a newline has been written.
   ///
   /// If this occurs in a place where newlines are prohibited, then invalidates
   /// the solution.
   void _handleNewline() {
-    if (!_options.allowNewlines) _hasInvalidNewline = true;
+    if (!_options.last.allowNewlines) _solution.invalidate(_options.last.piece);
 
     // Note that this piece contains a newline so that we can propagate that
     // up to containing pieces too.
-    _options.hasNewline = true;
+    _options.last.hasNewline = true;
   }
 
   /// Write any pending whitespace.
@@ -314,15 +283,15 @@ class CodeWriter {
   void _finishLine() {
     // If the completed line is too long, track the overflow.
     if (_column >= _pageWidth) {
-      _overflow += _column - _pageWidth;
+      _solution.addOverflow(_column - _pageWidth);
     }
 
     // If we found a problematic line, and there is a piece on the line that
     // we can try to split, then remember that piece so that the solution will
     // expand it next.
     if (!_foundExpandLine &&
         _nextPieceToExpand != null &&
-        (_column > _pageWidth || _hasInvalidNewline)) {
+        (_column > _pageWidth || !_solution.isValid)) {
       // We found a problematic line, so remember it and the piece on it.
       _foundExpandLine = true;
     } else if (!_foundExpandLine) {
@@ -359,6 +328,9 @@ enum Whitespace {
 
 /// The mutable state local to a single piece being formatted.
 class _PieceOptions {
+  /// The piece being formatted with these options.
+  final Piece piece;
+
   /// The absolute number of spaces of leading indentation coming from
   /// block-like structure or explicit extra indentation (aligning constructor
   /// initializers, `show` clauses, etc.).
@@ -373,5 +345,5 @@ class _PieceOptions {
   /// Whether any newlines have occurred in this piece or any of its children.
   bool hasNewline = false;
 
-  _PieceOptions(this.indent, this.allowNewlines);
+  _PieceOptions(this.piece, this.indent, this.allowNewlines);
 }