Refactor how piece constraints are applied. (#1343)

Refactor how piece constraints are applied.

When I added support for constructors, I introduced a little API to let
pieces enforce constraints between their state and the states of other
child pieces. It worked but felt a little bolted on. In particular,
enforcing the constraints during CodeWriter formatting meant we didn't
prune invalid solutions as efficiently as we could.

I'm working on if elements now and I needed a better way to handle
constraints between pieces for that, so I went ahead and refactored the
way this works to be, I think, generally better:

- Enforce constraints between pieces during solution expansion instead
  of during formatting. This means that we eagerly discard any solution
  (and the whole tree of solutions that could have come from it) if it
  doesn't meet the constraints.

- Remove the `_pieces` fields in PieceStateSet. With the previous
  optimization to track the next unsolved piece in Solution, that list
  wasn't doing much of anything. It's only use (aside from debug output)
  was for determining the order that pieces are compared when comparing
  two solutions. But we can get that from iterating over the
  _pieceStates map since maps in Dart track their key insertion order.

- Update ConstructorPiece to use this new API.

Co-authored-by: Nate Bosch <[email protected]>

Author: munificent

lib/src/back_end/code_writer.dart

@@ -53,16 +53,9 @@ class CodeWriter {
   /// The number of characters in the line currently being written.
   int _column = 0;
 
-  /// Whether this solution has encountered a conflict that makes it not a valid
-  /// solution.
-  ///
-  /// This can occur if:
-  ///
-  /// - A mandatory newline occurs like from a line comment or statement where
-  ///   none is permitted.
-  /// - An outer piece requires a child piece to have a certain state and it
-  ///   doesn't.
-  bool _isInvalid = false;
+  /// 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.
   ///
@@ -132,7 +125,7 @@ class CodeWriter {
 
     return Solution(_pieceStates, _buffer.toString(), _selectionStart,
         _selectionEnd, _nextPieceToExpand,
-        isValid: !_isInvalid, overflow: _overflow, cost: _cost);
+        isValid: !_hasInvalidNewline, overflow: _overflow, cost: _cost);
   }
 
   /// Notes that a newline has been written.
@@ -144,7 +137,7 @@ class CodeWriter {
   /// the raw text contains a newline, which can happen in multi-line block
   /// comments and multi-line string literals.
   void handleNewline() {
-    if (!_options.allowNewlines) _isInvalid = true;
+    if (!_options.allowNewlines) _hasInvalidNewline = true;
 
     // Note that this piece contains a newline so that we can propagate that
     // up to containing pieces too.
@@ -234,32 +227,13 @@ class CodeWriter {
 
   /// Format [piece] and insert the result into the code being written and
   /// returned by [finish()].
-  ///
-  /// If [requireState] is given, then [piece] must be in that state or the
-  /// solution is considered invalid. This is used to create constraints
-  /// between pieces. For example, a constructor piece forces the initializer
-  /// list child piece to split if the parameter list splits.
-  void format(Piece piece, {State? requireState}) {
+  void format(Piece piece) {
     _pieceOptions.add(_PieceOptions(_options.indent, _options.allowNewlines));
 
     var isUnsolved = !_pieceStates.isBound(piece) && piece.states.length > 1;
     if (isUnsolved) _currentUnsolvedPieces.add(piece);
 
     var state = _pieceStates.pieceState(piece);
-
-    // If the parent piece constrains this child to a certain state, then
-    // invalidate the solution if it doesn't match.
-    if (requireState != null && state != requireState) {
-      _isInvalid = true;
-      // TODO(perf): The solver doesn't discard invalid states that are caused
-      // by newlines because sometimes an invalid state like that can lead to
-      // a valid state by selecting values for other pieces.
-      //
-      // But in this case, once a partial solution is invalidated by one piece's
-      // state conflicting with another's, any solution derived from that is
-      // also going to be invalid and we can prune that entire solution tree.
-    }
-
     _cost += piece.stateCost(state);
 
     // TODO(perf): Memoize this. Might want to create a nested PieceWriter
@@ -336,7 +310,7 @@ class CodeWriter {
     // expand it next.
     if (!_foundExpandLine &&
         _nextPieceToExpand != null &&
-        (_column > _pageWidth || _isInvalid)) {
+        (_column > _pageWidth || _hasInvalidNewline)) {
       // We found a problematic line, so remember it and the piece on it.
       _foundExpandLine = true;
     } else if (!_foundExpandLine) {

lib/src/back_end/solution.dart

@@ -7,23 +7,19 @@ import 'code_writer.dart';
 /// A possibly incomplete set of selected states for a set of pieces being
 /// solved.
 class PieceStateSet {
-  // TODO(perf): Looking up and expanding the set of chunk states was a
-  // performance bottleneck in the old line splitter. If that turns out to be
-  // true here, then consider a faster representation for this list and the
-  // subsequent map field.
-  /// The in-order flattened list of all pieces being solved.
+  /// The states that pieces have been bound to.
   ///
-  /// This doesn't include pieces like text that have only a single value since
-  /// there's nothing to solve for them.
-  final List<Piece> _pieces;
-
+  /// Note that order that keys are inserted into this map is significant. When
+  /// ordering solutions, we use the order that pieces are bound in here to
+  /// break ties between solutions that otherwise have the same cost and
+  /// overflow.
   final Map<Piece, State> _pieceStates;
 
   /// Creates a new [PieceStateSet] with no pieces set to any state (which
   /// implicitly means they have state 0).
-  PieceStateSet(this._pieces) : _pieceStates = {};
+  PieceStateSet() : _pieceStates = {};
 
-  PieceStateSet._(this._pieces, this._pieceStates);
+  PieceStateSet._(this._pieceStates);
 
   /// The state this solution selects for [piece].
   ///
@@ -33,18 +29,43 @@ class PieceStateSet {
   /// Whether [piece] has been bound to a state in this set.
   bool isBound(Piece piece) => _pieceStates.containsKey(piece);
 
-  /// Creates a clone of this state with [piece] bound to [state].
-  PieceStateSet cloneWith(Piece piece, State state) {
-    return PieceStateSet._(_pieces, {..._pieceStates, piece: state});
+  /// Attempts to bind [piece] to [state], taking into account any constraints
+  /// pieces place on each other.
+  ///
+  /// Returns a new [PieceStateSet] with [piece] bound to [state] and any other
+  /// pieces constrained by that choice bound to their constrained values
+  /// (recursively). Returns `null` if a constraint conflicts with the already
+  /// bound or pinned state for some piece.
+  PieceStateSet? tryBind(Piece piece, State state) {
+    var conflict = false;
+    var boundStates = {..._pieceStates};
+
+    void traverse(Piece thisPiece, State thisState) {
+      // If this piece is already pinned or bound to some other state, then the
+      // solution doesn't make sense.
+      var alreadyBound = thisPiece.pinnedState ?? boundStates[thisPiece];
+      if (alreadyBound != null && alreadyBound != thisState) {
+        conflict = true;
+        return;
+      }
+
+      boundStates[thisPiece] = thisState;
+
+      // This piece may in turn place further constraints on others.
+      thisPiece.applyConstraints(thisState, traverse);
+    }
+
+    traverse(piece, state);
+
+    if (conflict) return null;
+    return PieceStateSet._(boundStates);
   }
 
   @override
   String toString() {
-    return _pieces.map((piece) {
-      var state = _pieceStates[piece];
-      var stateLabel = state == null ? '?' : '$state';
-      return '$piece:$stateLabel';
-    }).join(' ');
+    return _pieceStates.keys
+        .map((piece) => '$piece:${_pieceStates[piece]}')
+        .join(' ');
   }
 }
 
@@ -54,7 +75,7 @@ class PieceStateSet {
 /// code and its cost.
 class Solution implements Comparable<Solution> {
   /// The states the pieces have been set to in this solution.
-  final PieceStateSet _state;
+  final PieceStateSet _stateSet;
 
   /// The formatted code.
   final String text;
@@ -119,8 +140,8 @@ class Solution implements Comparable<Solution> {
   /// no selection.
   final int? selectionEnd;
 
-  factory Solution.initial(Piece root, int pageWidth, List<Piece> pieces) {
-    return Solution._(root, pageWidth, PieceStateSet(pieces));
+  factory Solution.initial(Piece root, int pageWidth) {
+    return Solution._(root, pageWidth, PieceStateSet());
   }
 
   factory Solution._(Piece root, int pageWidth, PieceStateSet state) {
@@ -129,7 +150,7 @@ class Solution implements Comparable<Solution> {
     return writer.finish();
   }
 
-  Solution(this._state, this.text, this.selectionStart, this.selectionEnd,
+  Solution(this._stateSet, this.text, this.selectionStart, this.selectionEnd,
       this._nextPieceToExpand,
       {required this.overflow, required this.cost, required this.isValid});
 
@@ -139,7 +160,8 @@ class Solution implements Comparable<Solution> {
     if (_nextPieceToExpand case var piece?) {
       return [
         for (var state in piece.states)
-          Solution._(root, pageWidth, _state.cloneWith(piece, state))
+          if (_stateSet.tryBind(piece, state) case final stateSet?)
+            Solution._(root, pageWidth, stateSet)
       ];
     }
 
@@ -161,16 +183,10 @@ class Solution implements Comparable<Solution> {
 
     if (overflow != other.overflow) return overflow.compareTo(other.overflow);
 
-    // Should be solving the same set of pieces.
-    assert(_state._pieces.length == other._state._pieces.length);
-
-    // If all else is equal, prefer lower states in earlier pieces.
-    // TODO(tall): This might not be needed once piece scoring is more
-    // sophisticated.
-    for (var i = 0; i < _state._pieces.length; i++) {
-      var piece = _state._pieces[i];
-      var thisState = _state.pieceState(piece);
-      var otherState = other._state.pieceState(piece);
+    // If all else is equal, prefer lower states in earlier bound pieces.
+    for (var piece in _stateSet._pieceStates.keys) {
+      var thisState = _stateSet.pieceState(piece);
+      var otherState = other._stateSet.pieceState(piece);
       if (thisState != otherState) return thisState.compareTo(otherState);
     }
 
@@ -183,7 +199,7 @@ class Solution implements Comparable<Solution> {
       '\$$cost',
       if (overflow > 0) '($overflow over)',
       if (!isValid) '(invalid)',
-      '$_state',
+      '$_stateSet',
     ].join(' ');
   }
 }

lib/src/back_end/solver.dart

@@ -34,27 +34,10 @@ class Solver {
 
   Solver(this._pageWidth);
 
-  /// Finds the best set of line splits for [piece] and returns the resulting
-  /// formatted code.
-  Solution format(Piece piece) {
-    // Collect all of the pieces with states that can be selected.
-    var unsolvedPieces = <Piece>[];
-
-    void traverse(Piece piece) {
-      // We don't need to worry about selecting pieces that have only one state.
-      if (piece.states.length > 1) unsolvedPieces.add(piece);
-      piece.forEachChild(traverse);
-    }
-
-    traverse(piece);
-
-    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.initial(root, _pageWidth, pieces);
+  /// Finds the best set of line splits for [root] piece and returns the
+  /// resulting formatted code.
+  Solution format(Piece root) {
+    var solution = Solution.initial(root, _pageWidth);
     _queue.add(solution);
 
     // The lowest cost solution found so far that does overflow.

lib/src/piece/constructor.dart

@@ -107,30 +107,42 @@ class ConstructorPiece extends Piece {
           _splitBetweenInitializers
       ];
 
+  /// Apply constraints between how the parameters may split and how the
+  /// initializers may split.
   @override
-  void format(CodeWriter writer, State state) {
-    // There are constraints between how the parameters may split and now the
-    // initializers may split.
-    var (parameterState, initializerState) = switch (state) {
-      // If there are no initializers, the parameters can do whatever.
-      State.unsplit when _initializers == null => (null, null),
-      // All parameters and initializers on one line.
-      State.unsplit => (State.unsplit, State.unsplit),
-      // If the `:` splits, then the parameters can't.
-      _splitBeforeInitializers => (State.unsplit, State.split),
-      // The `) :` on its own line.
-      _splitBetweenInitializers => (State.split, State.split),
-      _ => throw ArgumentError(),
-    };
+  void applyConstraints(State state, Constrain constrain) {
+    // If there are no initializers, the parameters can do whatever.
+    if (_initializers case var initializers?) {
+      switch (state) {
+        case State.unsplit:
+          // All parameters and initializers on one line.
+          constrain(_parameters, State.unsplit);
+          constrain(initializers, State.unsplit);
+
+        case _splitBeforeInitializers:
+          // Only split before the `:` when the parameters fit on one line.
+          constrain(_parameters, State.unsplit);
+          constrain(initializers, State.split);
+
+        case _splitBetweenInitializers:
+          // Split both the parameters and initializers and put the `) :` on
+          // its own line.
+          constrain(_parameters, State.split);
+          constrain(initializers, State.split);
+      }
+    }
+  }
 
+  @override
+  void format(CodeWriter writer, State state) {
     // If there's a newline in the header or parameters (like a line comment
     // after the `)`), then don't allow the initializers to remain unsplit.
     if (_initializers != null && state == State.unsplit) {
       writer.setAllowNewlines(false);
     }
 
     writer.format(_header);
-    writer.format(_parameters, requireState: parameterState);
+    writer.format(_parameters);
 
     if (_redirect case var redirect?) {
       writer.space();
@@ -153,7 +165,7 @@ class ConstructorPiece extends Piece {
         writer.setIndent(4);
       }
 
-      writer.format(initializers, requireState: initializerState);
+      writer.format(initializers);
     }
 
     writer.setIndent(0);

lib/src/piece/piece.dart

@@ -4,6 +4,8 @@
 
 import '../back_end/code_writer.dart';
 
+typedef Constrain = void Function(Piece other, State constrainedState);
+
 /// Base class for the formatter's internal representation used for line
 /// splitting.
 ///
@@ -48,6 +50,14 @@ abstract class Piece {
   State? get pinnedState => _pinnedState;
   State? _pinnedState;
 
+  /// Apply any constraints that this piece places on other pieces when this
+  /// piece is bound to [state].
+  ///
+  /// A piece class can override this. For any child piece that it wants to
+  /// constrain when this piece is in [state], call [constrain] and pass in the
+  /// child piece and the state that child should be constrained to.
+  void applyConstraints(State state, Constrain constrain) {}
+
   /// Given that this piece is in [state], use [writer] to produce its formatted
   /// output.
   void format(CodeWriter writer, State state);