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.
diff --git a/lib/src/back_end/code_writer.dart b/lib/src/back_end/code_writer.dart
index ebad84d..88ff261 100644
--- a/lib/src/back_end/code_writer.dart
+++ b/lib/src/back_end/code_writer.dart
@@ -208,7 +208,8 @@
     _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
diff --git a/lib/src/back_end/solution.dart b/lib/src/back_end/solution.dart
index b2db3a6..9fa106b 100644
--- a/lib/src/back_end/solution.dart
+++ b/lib/src/back_end/solution.dart
@@ -97,7 +97,30 @@
   /// 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 @@
   /// 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 @@
     _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 @@
     // 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
-    ];
-  }
+    // 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};
 
-  /// 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;
+      var additionalCost = _tryBind(newStates, expandPiece, state);
 
-    void bind(Piece thisPiece, State thisState) {
-      // If we've already failed from a previous sibling's constraint violation,
-      // early out.
-      if (conflict) return;
+      // Discard the solution if we hit a constraint violation.
+      if (additionalCost == null) continue;
 
-      // 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;
-      }
-
-      newStates[thisPiece] = thisState;
-      newCost += thisPiece.stateCost(thisState);
-
-      // 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 @@
       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;
+  }
 }
diff --git a/lib/src/front_end/ast_node_visitor.dart b/lib/src/front_end/ast_node_visitor.dart
index fe81f68..5600b5b 100644
--- a/lib/src/front_end/ast_node_visitor.dart
+++ b/lib/src/front_end/ast_node_visitor.dart
@@ -857,7 +857,7 @@
 
   @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 @@
 
   @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 @@
 
     var body = nodePiece(node.body);
 
-    var piece = IfPiece();
+    var piece = IfPiece(isStatement: true);
     piece.add(condition, body, isBlock: node.body is Block);
     return piece;
   }
diff --git a/lib/src/piece/if.dart b/lib/src/piece/if.dart
index df2e238..c99a08e 100644
--- a/lib/src/piece/if.dart
+++ b/lib/src/piece/if.dart
@@ -10,8 +10,13 @@
 /// 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 @@
     // 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);
+        }
       }
     }
   }
diff --git a/lib/src/piece/piece.dart b/lib/src/piece/piece.dart
index dfa7550..b3bf8e0 100644
--- a/lib/src/piece/piece.dart
+++ b/lib/src/piece/piece.dart
@@ -16,15 +16,13 @@
 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.