Reorganize how Solution handles pinned states and constraints. (#1373)

Reorganize how Solution handles pinned states and constraints.

In the process of turning ListElement into a Piece, I ran into a couple
of bugs in how pinned states and constraints between pieces are handled.

The tests didn't catch those bugs because there are basically two bugs
that cancel out. The IfPiece was incorrectly pinning an empty else block
in an if statement to split. But the solver was incorrectly ignoring
pinned states in some cases. This fixes both bugs.

With this change, when generating a Solution, we eagerly bind as many
Pieces to States as we can when the Solution is first created before
running CodeWriter. To do that, we find all of the pinned Pieces and
bind them to their pinned States, then propagate any constraints from
that to other child Pieces.

This fixes a couple of bugs, but also, I think, leads to a cleaner
module where at Solution creation time, we fill in as many piece states
as we can based on what we know statically about the Piece tree (its
pinned states and their constraints). Then during solution exploring and
line splitting, we don't have to worry about pinned states at all.

This will, I think, align with future optimizations where we do some
other piece pinning before running the solver.

Author: munificent

lib/src/back_end/code_writer.dart

@@ -208,7 +208,8 @@ class CodeWriter {
     _options.add(_PieceOptions(piece, _options.lastOrNull?.indent ?? 0,
         _options.lastOrNull?.allowNewlines ?? true));
 
-    var isUnsolved = !_solution.isBound(piece) && piece.states.length > 1;
+    var isUnsolved =
+        !_solution.isBound(piece) && piece.additionalStates.isNotEmpty;
     if (isUnsolved) _currentUnsolvedPieces.add(piece);
 
     // TODO(perf): Memoize this. Might want to create a nested PieceWriter

lib/src/back_end/solution.dart

@@ -97,7 +97,30 @@ class Solution implements Comparable<Solution> {
   /// 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).
-  Solution(Piece root, int pageWidth) : this._(root, pageWidth, 0, {});
+  factory Solution(Piece root, int pageWidth) {
+    var pieceStates = <Piece, State>{};
+    var cost = 0;
+
+    // Bind every pinned piece to its state and propagate any constraints from
+    // those.
+    void traversePinned(Piece piece) {
+      if (piece.pinnedState case var pinned?) {
+        var additionalCost = _tryBind(pieceStates, piece, pinned);
+
+        // Pieces should be implemented such that they never get pinned into a
+        // conflicting state because then there's no possible solution, so
+        // [_tryBind()] should always succeed and [additionalCost] won't be
+        // `null`.
+        cost += additionalCost!;
+      }
+
+      piece.forEachChild(traversePinned);
+    }
+
+    traversePinned(root);
+
+    return Solution._(root, pageWidth, cost, pieceStates);
+  }
 
   Solution._(Piece root, int pageWidth, this.cost, this._pieceStates) {
     var writer = CodeWriter(pageWidth, this);
@@ -111,7 +134,7 @@ class Solution implements Comparable<Solution> {
   /// The state this solution selects for [piece].
   ///
   /// If no state has been selected, defaults to the first state.
-  State pieceState(Piece piece) => _pieceStates[piece] ?? piece.states.first;
+  State pieceState(Piece piece) => _pieceStates[piece] ?? State.unsplit;
 
   /// Whether [piece] has been bound to a state in this set.
   bool isBound(Piece piece) => _pieceStates.containsKey(piece);
@@ -147,8 +170,11 @@ class Solution implements Comparable<Solution> {
     _invalidPiece = piece;
   }
 
-  /// When called on a [Solution] with some unselected piece states, chooses a
-  /// piece and yields further solutions for each state that piece can have.
+  /// Derives new potential solutions from this one by binding
+  /// [_nextPieceToExpand] to all of its possible states.
+  ///
+  /// If there is no potential piece to expand, or all attempts to expand it
+  /// fail, returns an empty list.
   List<Solution> expand(Piece root, int pageWidth) {
     // If the piece whose newline constraint was violated is already bound to
     // one state, then every solution derived from this one will also fail in
@@ -176,51 +202,25 @@ class Solution implements Comparable<Solution> {
     // make any noticeable performance difference on the one pathological
     // example I tried. Leaving this here as a TODO to investigate more when
     // there are other benchmarks we can try.
+    var solutions = <Solution>[];
 
-    return [
-      for (var state in expandPiece.states)
-        if (_tryBind(root, pageWidth, expandPiece, state) case var solution?)
-          solution
-    ];
-  }
-
-  /// Attempts to extend this solution's piece states by binding [piece] to
-  /// [state], taking into account any constraints pieces place on each other.
-  ///
-  /// Returns a new [Solution] 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.
-  Solution? _tryBind(Piece root, int pageWidth, Piece piece, State state) {
-    var conflict = false;
-    var newStates = {..._pieceStates};
-    var newCost = cost;
+    // For each state that the expanding piece can be in, create a new solution
+    // that inherits all of the bindings of this one, and binds the expanding
+    // piece to that state (along with any further pieces constrained by that
+    // one).
+    for (var state in expandPiece.states) {
+      var newStates = {..._pieceStates};
 
-    void bind(Piece thisPiece, State thisState) {
-      // If we've already failed from a previous sibling's constraint violation,
-      // early out.
-      if (conflict) return;
-
-      // If this piece is already pinned or bound to some other state, then the
-      // solution doesn't make sense.
-      var alreadyBound = thisPiece.pinnedState ?? newStates[thisPiece];
-      if (alreadyBound != null && alreadyBound != thisState) {
-        conflict = true;
-        return;
-      }
+      var additionalCost = _tryBind(newStates, expandPiece, state);
 
-      newStates[thisPiece] = thisState;
-      newCost += thisPiece.stateCost(thisState);
+      // Discard the solution if we hit a constraint violation.
+      if (additionalCost == null) continue;
 
-      // This piece may in turn place further constraints on others.
-      thisPiece.applyConstraints(thisState, bind);
+      solutions
+          .add(Solution._(root, pageWidth, cost + additionalCost, newStates));
     }
 
-    bind(piece, state);
-
-    if (conflict) return null;
-
-    return Solution._(root, pageWidth, newCost, newStates);
+    return solutions;
   }
 
   /// Compares two solutions where a more desirable solution comes first.
@@ -260,4 +260,46 @@ class Solution implements Comparable<Solution> {
       states,
     ].join(' ');
   }
+
+  /// Attempts to add a binding from [piece] to [state] in [boundStates], and
+  /// then adds any further bindings from constraints that [piece] applies to
+  /// its children, recursively.
+  ///
+  /// This may fail if [piece] is already bound to a different [state], or if
+  /// any constrained pieces are bound to different states.
+  ///
+  /// If successful, returns the additional cost required to bind [piece] to
+  /// [state] (along with any other applied constrained pieces). Otherwise,
+  /// returns `null` to indicate failure.
+  static int? _tryBind(
+      Map<Piece, State> boundStates, Piece piece, State state) {
+    var success = true;
+    var additionalCost = 0;
+
+    void bind(Piece thisPiece, State thisState) {
+      // If we've already failed from a previous sibling's constraint violation,
+      // early out.
+      if (!success) return;
+
+      // Apply the new binding if it doesn't conflict with an existing one.
+      switch (boundStates[thisPiece]) {
+        case null:
+          // Binding a unbound piece to a state.
+          additionalCost += thisPiece.stateCost(thisState);
+          boundStates[thisPiece] = thisState;
+
+          // This piece may in turn place further constraints on others.
+          thisPiece.applyConstraints(thisState, bind);
+        case var alreadyBound when alreadyBound != thisState:
+          // Already bound to a different state, so there's a conflict.
+          success = false;
+        default:
+          break; // Already bound to the same state, so nothing to do.
+      }
+    }
+
+    bind(piece, state);
+
+    return success ? additionalCost : null;
+  }
 }

lib/src/front_end/ast_node_visitor.dart

@@ -857,7 +857,7 @@ class AstNodeVisitor extends ThrowingAstVisitor<Piece> with PieceFactory {
 
   @override
   Piece visitIfElement(IfElement node) {
-    var piece = IfPiece();
+    var piece = IfPiece(isStatement: false);
 
     // Recurses through the else branches to flatten them into a linear if-else
     // chain handled by a single [IfPiece].
@@ -940,7 +940,7 @@ class AstNodeVisitor extends ThrowingAstVisitor<Piece> with PieceFactory {
 
   @override
   Piece visitIfStatement(IfStatement node) {
-    var piece = IfPiece();
+    var piece = IfPiece(isStatement: true);
 
     // Recurses through the else branches to flatten them into a linear if-else
     // chain handled by a single [IfPiece].
@@ -1840,7 +1840,7 @@ class AstNodeVisitor extends ThrowingAstVisitor<Piece> with PieceFactory {
 
     var body = nodePiece(node.body);
 
-    var piece = IfPiece();
+    var piece = IfPiece(isStatement: true);
     piece.add(condition, body, isBlock: node.body is Block);
     return piece;
   }

lib/src/piece/if.dart

@@ -10,8 +10,13 @@ import 'piece.dart';
 /// We also use this for while statements, which are formatted exactly like an
 /// if statement with no else clause.
 class IfPiece extends Piece {
+  /// Whether this is an if statement versus if collection element.
+  final bool _isStatement;
+
   final List<_IfSection> _sections = [];
 
+  IfPiece({required bool isStatement}) : _isStatement = isStatement;
+
   void add(Piece header, Piece statement, {required bool isBlock}) {
     _sections.add(_IfSection(header, statement, isBlock));
   }
@@ -24,11 +29,12 @@ class IfPiece extends Piece {
     // If an if element, any spread collection's split state must follow the
     // surrounding if element's: we either split all the spreads or none of
     // them. And if any of the non-spread then or else branches split, then the
-    // spreads do too. This has no effect on if statements since blocks always
-    // split.
-    for (var section in _sections) {
-      if (section.isBlock) {
-        constrain(section.statement, state);
+    // spreads do too.
+    if (!_isStatement) {
+      for (var section in _sections) {
+        if (section.isBlock) {
+          constrain(section.statement, state);
+        }
       }
     }
   }

lib/src/piece/piece.dart

@@ -16,15 +16,13 @@ typedef Constrain = void Function(Piece other, State constrainedState);
 abstract class Piece {
   /// The ordered list of ways this piece may split.
   ///
-  /// If the piece is aready pinned, then this is just the one pinned state.
-  /// Otherwise, it's [State.unsplit], which all pieces support, followed by
-  /// any other [additionalStates].
+  /// This is [State.unsplit], which all pieces support, followed by any other
+  /// [additionalStates].
   List<State> get states {
-    if (_pinnedState case var state?) {
-      return [state];
-    } else {
-      return [State.unsplit, ...additionalStates];
-    }
+    // Pinned pieces should be bound eagerly in the Solution so we shouldn't
+    // ever need to iterate over their states.
+    assert(_pinnedState == null);
+    return [State.unsplit, ...additionalStates];
   }
 
   /// The ordered list of all possible ways this piece could split.