Invalidate eagerly by newline constraint (#1495)

Propagate newline constrains eagerly when binding states in a Solution.

One of the main ways the formatter defines its formatting style is by having constraints that prohibit newlines in some child pieces when a parent is in a given state.

The two simplest examples (which cover a large amount of code) is that when an InfixPiece or a ListPiece is in State.unsplit, they don't allow any newlines in any of their children. That constraint effectively creates the desired style which is that a split inside an infix operand or list element forces the surrounding expression to split.

Whenever we bind a parent piece to some state, we can now ask it if it allows any of its children to contain newlines. For any given child, if the answer is "no" then:

-   If the child is already bound to a state that contains newlines, then we know this solution is a dead end. It and every solution you can ever derive from it will contain an invalid newline.

-   If the child isn't bound to a state yet, we can still look at all of its states and see which of them contain newlines. Any state that does can be eliminated because we'll never successfully bind the child to that state without violating this constraint.

    It may turn out that no states are left, in which case again we have a dead end solution.

    Or there may be just one valid state (usually State.unsplit), and we can immediately the child to that state and do this whole process recursively for this child.

    Or there may be just a couple of states left and we can at least winnow this child's states down to that list when we go to expand it later.

The end result is that even before actually formatting a solution, we can often tell if it's a dead end and discard it. If not, we can often bind a whole bunch of the bound piece's children in one go instead of having to do them one at a time and slowly formatting the interim solutions at each step.

The end result is a pretty decent improvement. Here's the micro benchmarks:

```
Benchmark (tall)                fastest   median  slowest  average  baseline
-----------------------------  --------  -------  -------  -------  --------
block                             0.070    0.071    0.122    0.075     92.7%
chain                             0.640    0.654    0.681    0.655    254.5%
collection                        0.175    0.180    0.193    0.181     96.3%
collection_large                  0.930    0.955    0.988    0.955     97.2%
conditional                       0.067    0.068    0.086    0.071    132.4%
curry                             0.596    0.608    1.478    0.631    278.9%
flutter_popup_menu_test           0.293    0.302    0.326    0.303    141.4%
flutter_scrollbar_test            0.160    0.166    0.184    0.166     96.1%
function_call                     1.460    1.483    1.680    1.490     97.5%
infix_large                       0.709    0.733    0.761    0.733     97.4%
infix_small                       0.175    0.181    0.198    0.181     93.0%
interpolation                     0.091    0.096    0.118    0.096     98.7%
large                             3.514    3.574    3.767    3.596    129.5%
top_level                         0.146    0.150    0.182    0.152    106.7%
```

There's a slight regression on some of the tiny microbenchmarks (probably because there's some overhead traversing and binding children to State.unsplit when that solution ends up winning anyway). But the improvement to the larger ones is significant.

More interesting is the overall performance. Here's formatting the Flutter repo:

```
Current formatter     10.964 ========================================
This PR                8.826 ================================
Short style formatter  4.558 ================
```

The current formatter is 58.43% slower than the old short style formatter.

This PR is 24.22% faster than the current formatter. It's 48.36% slower than the old formatter, but it gets the formatter 33.37% of the way to the old short style one.
diff --git a/lib/src/back_end/code_writer.dart b/lib/src/back_end/code_writer.dart
index 443da6b..526f9af 100644
--- a/lib/src/back_end/code_writer.dart
+++ b/lib/src/back_end/code_writer.dart
@@ -209,7 +209,7 @@
   /// and multi-line strings.
   void whitespace(Whitespace whitespace, {bool flushLeft = false}) {
     if (whitespace case Whitespace.newline || Whitespace.blankLine) {
-      _handleNewline(allowNewlines: true);
+      _hadNewline = true;
       _pendingIndent = flushLeft ? 0 : _indentStack.last.indent;
     }
 
@@ -231,7 +231,7 @@
   /// be `false`. It's up to the parent piece to only call this when it's safe
   /// to do so. In practice, this usually means when the parent piece knows that
   /// [piece] will have a newline before and after it.
-  void format(Piece piece, {bool separate = false, bool allowNewlines = true}) {
+  void format(Piece piece, {bool separate = false}) {
     if (separate) {
       Profile.count('CodeWriter.format() piece separate');
 
@@ -239,7 +239,7 @@
     } else {
       Profile.count('CodeWriter.format() piece inline');
 
-      _formatInline(piece, allowNewlines: allowNewlines);
+      _formatInline(piece);
     }
   }
 
@@ -270,7 +270,7 @@
   }
 
   /// Format [piece] writing directly into this [CodeWriter].
-  void _formatInline(Piece piece, {required bool allowNewlines}) {
+  void _formatInline(Piece piece) {
     // Begin a new formatting context for this child.
     var previousPiece = _currentPiece;
     _currentPiece = piece;
@@ -309,9 +309,25 @@
 
     _currentPiece = previousPiece;
 
-    // If the child contained a newline then the parent transitively does.
-    if (childHadNewline && _currentPiece != null) {
-      _handleNewline(allowNewlines: allowNewlines);
+    // If the child contained a newline then invalidate the solution if any of
+    // the containing pieces don't allow one at this point in the tree.
+    if (childHadNewline) {
+      // TODO(rnystrom): We already do much of the newline constraint validation
+      // when the Solution is first created before we format. For performance,
+      // it would be good to do *all* of it before formatting. The missing part
+      // is that pieces containing hard newlines (comments, multiline strings,
+      // sequences, etc.) do not constrain their parents when the solution is
+      // first created. If we can get that working, then this check can be
+      // removed.
+      if (_currentPiece case var parent?
+          when !parent.allowNewlineInChild(
+              _solution.pieceState(parent), piece)) {
+        _solution.invalidate(_currentPiece!);
+      }
+
+      // Note that this piece contains a newline so that we can propagate that
+      // up to containing pieces too.
+      _hadNewline = true;
     }
   }
 
@@ -327,18 +343,6 @@
     _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({required bool allowNewlines}) {
-    if (!allowNewlines) _solution.invalidate(_currentPiece!);
-
-    // Note that this piece contains a newline so that we can propagate that
-    // up to containing pieces too.
-    _hadNewline = true;
-  }
-
   /// Write any pending whitespace.
   ///
   /// This is called before non-whitespace text is about to be written, or
diff --git a/lib/src/back_end/solution.dart b/lib/src/back_end/solution.dart
index 9dedbb7..69cba90 100644
--- a/lib/src/back_end/solution.dart
+++ b/lib/src/back_end/solution.dart
@@ -25,6 +25,15 @@
   /// overflow.
   final Map<Piece, State> _pieceStates;
 
+  /// The set of states that pieces are allowed to be in without violating
+  /// constraints of already bound pieces.
+  ///
+  /// Each key is a constrained piece and the values are the remaining states
+  /// that the piece may take which aren't known to violate existing
+  /// constraints. If a piece is not in this map, then there are no constraints
+  /// on it.
+  final Map<Piece, List<State>> _allowedStates;
+
   /// The amount of penalties applied based on the chosen line splits.
   int get cost => _cost;
   int _cost;
@@ -33,16 +42,19 @@
   String get text => _text;
   late final String _text;
 
-  /// Whether this score is for a valid solution or not.
+  /// False if this Solution contains a newline where one is prohibited.
   ///
-  /// An invalid solution is one where a hard newline appears in a context
-  /// where splitting isn't allowed. This is considered worse than any other
-  /// solution.
-  bool get isValid => _invalidPiece == null;
+  /// An invalid solution may have no overflow characters and the lowest score,
+  /// but is still considered worse than any other valid solution.
+  bool get isValid => _isValid;
+  bool _isValid = true;
 
-  /// The piece that forbid newlines when an invalid newline was written, or
-  /// `null` if no invalid newline has occurred.
-  Piece? _invalidPiece;
+  /// Whether the solution contains an invalid newline and the piece that
+  /// prohibits the newline is bound in this solution.
+  ///
+  /// When this is `true`, it means this solution and every solution that could
+  /// be derived from it is invalid so the whole solution tree can be discarded.
+  bool _isDeadEnd = false;
 
   /// The total number of characters that do not fit inside the page width.
   int get overflow => _overflow;
@@ -102,33 +114,20 @@
   /// another state).
   factory Solution(SolutionCache cache, Piece root,
       {required int pageWidth, required int leadingIndent, State? rootState}) {
-    var pieceStates = <Piece, State>{};
-    var cost = 0;
-
-    // If we're formatting a subtree of a larger Piece tree that binds [root]
-    // to [rootState], then bind it in this solution too.
-    if (rootState != null) {
-      var additionalCost = _tryBind(pieceStates, root, rootState);
-
-      // Binding should always succeed since we should only get here when
-      // formatting a subtree whose surrounding Solution successfully bound
-      // this piece to this state.
-      cost += additionalCost!;
-    }
-
-    return Solution._(cache, root, pageWidth, leadingIndent, cost, pieceStates);
+    var solution =
+        Solution._(cache, root, pageWidth, leadingIndent, 0, {}, {}, rootState);
+    solution._format(cache, root, pageWidth, leadingIndent);
+    return solution;
   }
 
   Solution._(SolutionCache cache, Piece root, int pageWidth, int leadingIndent,
-      this._cost, this._pieceStates) {
+      this._cost, this._pieceStates, this._allowedStates,
+      [State? rootState]) {
     Profile.count('create Solution');
 
-    var writer = CodeWriter(pageWidth, leadingIndent, cache, this);
-    writer.format(root);
-    var (text, expandPieces) = writer.finish();
-
-    _text = text;
-    _expandPieces = expandPieces;
+    // If we're formatting a subtree of a larger Piece tree that binds [root]
+    // to [rootState], then bind it in this solution too.
+    if (rootState != null) _bind(root, rootState);
   }
 
   /// Attempt to eagerly bind [piece] to a state given that it must fit within
@@ -139,8 +138,7 @@
   /// `true`. Otherwise returns `false`.
   bool tryBindByPageWidth(Piece piece, int pageWidth) {
     if (piece.fixedStateForPageWidth(pageWidth) case var state?) {
-      var cost = _tryBind(_pieceStates, piece, state);
-      _cost += cost!;
+      _bind(piece, state);
       return true;
     }
 
@@ -212,7 +210,11 @@
   ///
   /// This should only be called by [CodeWriter].
   void invalidate(Piece piece) {
-    _invalidPiece = piece;
+    _isValid = false;
+
+    // If the piece whose newline constraint was violated is already bound to
+    // one state, then every solution derived from this one will also fail.
+    if (!_isDeadEnd && isBound(piece)) _isDeadEnd = true;
   }
 
   /// Derives new potential solutions from this one by binding [_expandPieces]
@@ -222,11 +224,6 @@
   /// fail, returns an empty list.
   List<Solution> expand(SolutionCache cache, Piece root,
       {required int pageWidth, required int leadingIndent}) {
-    // 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
-    // the same way, so discard the whole solution tree hanging off this one.
-    if (_invalidPiece case var piece? when isBound(piece)) return const [];
-
     // If there is no piece that we can expand on this solution, it's a dead
     // end (or a winner).
     if (_expandPieces.isEmpty) return const [];
@@ -238,19 +235,19 @@
       // the expanding piece to that state (along with any further pieces
       // constrained by that one).
       var expandPiece = _expandPieces[i];
-      for (var state in expandPiece.additionalStates) {
-        var newStates = {..._pieceStates};
+      for (var state
+          in _allowedStates[expandPiece] ?? expandPiece.additionalStates) {
+        var expanded = Solution._(cache, root, pageWidth, leadingIndent, cost,
+            {..._pieceStates}, {..._allowedStates});
 
         // Bind all preceding expand pieces to their unsplit state. Their
         // other states have already been expanded by earlier iterations of
         // the outer for loop.
         var valid = true;
-        var additionalCost = 0;
         for (var j = 0; j < i; j++) {
-          if (_tryBind(newStates, _expandPieces[j], State.unsplit)
-              case var cost?) {
-            additionalCost += cost;
-          } else {
+          expanded._bind(_expandPieces[j], State.unsplit);
+
+          if (expanded._isDeadEnd) {
             valid = false;
             break;
           }
@@ -259,15 +256,19 @@
         // Discard the solution if we hit a constraint violation.
         if (!valid) continue;
 
-        if (_tryBind(newStates, expandPiece, state) case var cost?) {
-          additionalCost += cost;
-        } else {
-          // Discard the solution if we hit a constraint violation.
-          continue;
-        }
+        expanded._bind(expandPiece, state);
 
-        solutions.add(Solution._(cache, root, pageWidth, leadingIndent,
-            cost + additionalCost, newStates));
+        // Discard the solution if we hit a constraint violation.
+        if (!expanded._isDeadEnd) {
+          expanded._format(cache, root, pageWidth, leadingIndent);
+
+          // TODO(rnystrom): These come mostly (entirely?) from hard newlines
+          // in sequences, comments, and multiline strings. It should be
+          // possible to handle those during piece construction too. If we do,
+          // remove this check.
+          // We may not detect some newline violations until formatting.
+          if (!expanded._isDeadEnd) solutions.add(expanded);
+        }
       }
     }
 
@@ -314,45 +315,112 @@
     ].join(' ').trim();
   }
 
-  /// Attempts to add a binding from [piece] to [state] in [boundStates], and
+  /// Run a [CodeWriter] on this solution to produce the final formatted output
+  /// and calculate the overflow and expand pieces.
+  void _format(
+      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;
+    _expandPieces = expandPieces;
+  }
+
+  /// Attempts to add a binding from [piece] to [state] to the solution, 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.
+  /// This may invalidate the solution 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;
+  /// If successful, adds the cost required to bind [piece] to [state] (along
+  /// with any other applied constrained pieces). Otherwise, marks the solution
+  /// as a dead end.
+  void _bind(Piece piece, State state) {
+    // If we've already failed from a previous violation, early out.
+    if (_isDeadEnd) return;
 
-    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 (pieceStateIfBound(piece)) {
+      case null:
+        // Binding a unbound piece to a state.
+        _cost += piece.stateCost(state);
+        _pieceStates[piece] = state;
 
-      // Apply the new binding if it doesn't conflict with an existing one.
-      switch (thisPiece.pinnedState ?? 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.
+        piece.applyConstraints(state, _bind);
 
-          // 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.
+        // If this piece's state prevents some of its children from having
+        // newlines, then further constrain those children.
+        if (!_isDeadEnd) {
+          piece.forEachChild((child) {
+            // Stop as soon as we fail.
+            if (_isDeadEnd) return;
+
+            // If the child can have newlines, there is no constraint.
+            if (piece.allowNewlineInChild(state, child)) return;
+
+            // Otherwise, don't let any piece under [child] contain newlines.
+            _constrainOffspring(child);
+          });
+        }
+
+      case var alreadyBound when alreadyBound != state:
+        // Already bound to a different state, so there's a conflict.
+        _isDeadEnd = true;
+        _isValid = false;
+
+      default:
+        break; // Already bound to the same state, so nothing to do.
+    }
+  }
+
+  /// For [piece] and its transitive offspring subtree, eliminate any state that
+  /// will always produce a newline since that state is not permitted because
+  /// the parent of [piece] doesn't allow [piece] to have any newlines.
+  void _constrainOffspring(Piece piece) {
+    for (var offspring in piece.statefulOffspring) {
+      if (_isDeadEnd) break;
+
+      if (pieceStateIfBound(offspring) case var boundState?) {
+        // This offspring is already pinned or bound to a state. If that
+        // state will emit newlines, then this solution is invalid.
+        if (offspring.containsNewline(boundState)) {
+          _isDeadEnd = true;
+          _isValid = false;
+        }
+      } else if (!_allowedStates.containsKey(offspring)) {
+        // If we get here, the offspring isn't bound to a state and we haven't
+        // already constrained it. Eliminate any of its states that will emit
+        // newlines.
+        var allowedUnsplit = !offspring.containsNewline(State.unsplit);
+
+        var states = offspring.additionalStates;
+        var remainingStates = <State>[];
+        for (var state in states) {
+          if (!offspring.containsNewline(state)) {
+            remainingStates.add(state);
+          }
+        }
+
+        if (!allowedUnsplit && remainingStates.isEmpty) {
+          // There is no state this child can take that won't emit newlines,
+          // and it's not allowed to, so this solution is bad.
+          _isDeadEnd = true;
+          _isValid = false;
+        } else if (remainingStates.isEmpty) {
+          // The only valid state is unsplit so bind it to that.
+          _bind(offspring, State.unsplit);
+        } else if (!allowedUnsplit && remainingStates.length == 1) {
+          // There's only one valid state, so bind it to that.
+          _bind(offspring, remainingStates.first);
+        } else if (remainingStates.length < states.length) {
+          // There are some constrained states, so keep the remaining ones.
+          _allowedStates[offspring] = remainingStates;
+        }
       }
     }
-
-    bind(piece, state);
-
-    return success ? additionalCost : null;
   }
 }
diff --git a/lib/src/front_end/piece_writer.dart b/lib/src/front_end/piece_writer.dart
index a70f3e7..7ccc138 100644
--- a/lib/src/front_end/piece_writer.dart
+++ b/lib/src/front_end/piece_writer.dart
@@ -12,6 +12,7 @@
 import '../piece/adjacent.dart';
 import '../piece/list.dart';
 import '../piece/piece.dart';
+import '../piece/text.dart';
 import '../profile.dart';
 import '../source_code.dart';
 import 'comment_writer.dart';
diff --git a/lib/src/front_end/sequence_builder.dart b/lib/src/front_end/sequence_builder.dart
index 7d16f55..3c7cf39 100644
--- a/lib/src/front_end/sequence_builder.dart
+++ b/lib/src/front_end/sequence_builder.dart
@@ -8,6 +8,7 @@
 import '../constants.dart';
 import '../piece/piece.dart';
 import '../piece/sequence.dart';
+import '../piece/text.dart';
 import 'piece_factory.dart';
 
 /// Incrementally builds a [SequencePiece], including handling comments and
@@ -38,9 +39,6 @@
 
   SequenceBuilder(this._visitor);
 
-  bool _mustSplit = false;
-  bool get mustSplit => _mustSplit;
-
   Piece build({bool forceSplit = false}) {
     // If the sequence only contains a single piece, just return it directly
     // and discard the unnecessary wrapping.
@@ -51,17 +49,27 @@
       return _elements.single.piece;
     }
 
-    // Discard any trailing blank line after the last element.
-    if (_elements.isNotEmpty) {
-      _elements.last.blankAfter = false;
+    // If there are no elements, don't bother making a SequencePiece or
+    // BlockPiece.
+    if (_elements.isEmpty) {
+      return _visitor.pieces.build(() {
+        _visitor.pieces.add(_leftBracket!);
+        if (forceSplit) {
+          _visitor.pieces.add(NewlinePiece());
+        }
+        _visitor.pieces.add(_rightBracket!);
+      });
     }
 
-    var piece = SequencePiece(_elements,
-        leftBracket: _leftBracket, rightBracket: _rightBracket);
+    // Discard any trailing blank line after the last element.
+    _elements.last.blankAfter = false;
 
-    if (mustSplit || forceSplit) piece.pin(State.split);
+    var sequence = SequencePiece(_elements);
+    if ((_leftBracket, _rightBracket) case (var left?, var right?)) {
+      return BlockPiece(left, sequence, right);
+    }
 
-    return piece;
+    return sequence;
   }
 
   /// Adds the opening [bracket] to the built sequence.
@@ -82,12 +90,9 @@
   /// The caller should have already called [addCommentsBefore()] with the
   /// first token in [piece].
   void add(Piece piece, {int? indent, bool allowBlankAfter = true}) {
-    _add(indent ?? Indent.none, piece);
+    _elements.add(SequenceElementPiece(indent ?? Indent.none, piece));
 
     _allowBlank = allowBlankAfter;
-
-    // There is at least one non-comment element in the sequence, so it splits.
-    _mustSplit = true;
   }
 
   /// Visits [node] and adds the resulting [Piece] to this sequence, handling
@@ -147,13 +152,10 @@
         }
 
         // Write the comment as its own sequence piece.
-        _add(indent, comment);
+        _elements.add(SequenceElementPiece(indent, comment));
       }
     }
 
-    // If the sequence contains any line comments, make sure it splits.
-    if (comments.requiresNewline) _mustSplit = true;
-
     // Write a blank before the token if there should be one.
     if (comments.linesBeforeNextToken > 1) {
       // If we just wrote a comment, then allow a blank line between it and the
@@ -163,11 +165,4 @@
       addBlank();
     }
   }
-
-  void _add(int indent, Piece piece) {
-    // If there are brackets, add a level of block indentation.
-    if (_leftBracket != null) indent += Indent.block;
-
-    _elements.add(SequenceElementPiece(indent, piece));
-  }
 }
diff --git a/lib/src/piece/assign.dart b/lib/src/piece/assign.dart
index 8aa9412..40f473c 100644
--- a/lib/src/piece/assign.dart
+++ b/lib/src/piece/assign.dart
@@ -140,6 +140,18 @@
   }
 
   @override
+  bool allowNewlineInChild(State state, Piece child) {
+    if (state == State.unsplit) {
+      if (child == _left) return false;
+
+      // Always allow block-splitting the right side if it supports it.
+      if (child == _right) return _canBlockSplitRight;
+    }
+
+    return true;
+  }
+
+  @override
   void format(CodeWriter writer, State state) {
     switch (state) {
       case State.unsplit:
@@ -170,7 +182,7 @@
   }
 
   void _writeLeft(CodeWriter writer, {bool allowNewlines = true}) {
-    if (_left case var left?) writer.format(left, allowNewlines: allowNewlines);
+    if (_left case var left?) writer.format(left);
   }
 
   void _writeOperator(CodeWriter writer, {bool split = false}) {
@@ -183,7 +195,7 @@
   void _writeRight(CodeWriter writer,
       {bool indent = false, bool allowNewlines = true}) {
     if (indent) writer.pushIndent(Indent.expression);
-    writer.format(_right, allowNewlines: allowNewlines);
+    writer.format(_right);
     if (indent) writer.popIndent();
   }
 
@@ -217,7 +229,7 @@
     // can split because there are no block operands.
     var totalLength = 0;
     if (_left case var left? when !_canBlockSplitLeft) {
-      if (left.containsNewline) return _atOperator;
+      if (left.containsHardNewline) return _atOperator;
 
       totalLength += left.totalCharacters;
     }
@@ -225,7 +237,7 @@
     totalLength += _operator.totalCharacters;
 
     if (!_canBlockSplitRight) {
-      if (_right.containsNewline) return _atOperator;
+      if (_right.containsHardNewline) return _atOperator;
       totalLength += _right.totalCharacters;
     }
 
diff --git a/lib/src/piece/case.dart b/lib/src/piece/case.dart
index a4a8557..de42ef3 100644
--- a/lib/src/piece/case.dart
+++ b/lib/src/piece/case.dart
@@ -59,25 +59,20 @@
       ];
 
   @override
+  bool allowNewlineInChild(State state, Piece child) {
+    return switch (state) {
+      _ when child == _arrow => true,
+      State.unsplit when child == _body && _canBlockSplitBody => true,
+      _beforeBody when child == _pattern =>
+        _guard == null || _patternIsLogicalOr,
+      _beforeBody when child == _body => true,
+      _beforeWhenAndBody => true,
+      _ => false,
+    };
+  }
+
+  @override
   void format(CodeWriter writer, State state) {
-    var allowNewlineInPattern = false;
-    var allowNewlineInGuard = false;
-    var allowNewlineInBody = false;
-
-    switch (state) {
-      case State.unsplit:
-        allowNewlineInBody = _canBlockSplitBody;
-
-      case _beforeBody:
-        allowNewlineInPattern = _guard == null || _patternIsLogicalOr;
-        allowNewlineInBody = true;
-
-      case _beforeWhenAndBody:
-        allowNewlineInPattern = true;
-        allowNewlineInGuard = true;
-        allowNewlineInBody = true;
-    }
-
     // If there is a split guard, then indent the pattern past it.
     var indentPatternForGuard = !_canBlockSplitPattern &&
         !_patternIsLogicalOr &&
@@ -85,14 +80,14 @@
 
     if (indentPatternForGuard) writer.pushIndent(Indent.expression);
 
-    writer.format(_pattern, allowNewlines: allowNewlineInPattern);
+    writer.format(_pattern);
 
     if (indentPatternForGuard) writer.popIndent();
 
     if (_guard case var guard?) {
       writer.pushIndent(Indent.expression);
       writer.splitIf(state == _beforeWhenAndBody);
-      writer.format(guard, allowNewlines: allowNewlineInGuard);
+      writer.format(guard);
       writer.popIndent();
     }
 
@@ -102,7 +97,7 @@
     if (state != State.unsplit) writer.pushIndent(Indent.block);
 
     writer.splitIf(state == _beforeBody || state == _beforeWhenAndBody);
-    writer.format(_body, allowNewlines: allowNewlineInBody);
+    writer.format(_body);
 
     if (state != State.unsplit) writer.popIndent();
   }
diff --git a/lib/src/piece/chain.dart b/lib/src/piece/chain.dart
index 55497f7..b7a2d38 100644
--- a/lib/src/piece/chain.dart
+++ b/lib/src/piece/chain.dart
@@ -161,10 +161,31 @@
   }
 
   @override
+  bool allowNewlineInChild(State state, Piece child) {
+    switch (state) {
+      case _ when child == _target:
+        return _allowSplitInTarget || state == State.split;
+
+      case State.unsplit:
+        return false;
+
+      case _splitAfterProperties:
+        for (var i = 0; i < _leadingProperties; i++) {
+          if (_calls[i]._call == child) return false;
+        }
+
+      case _blockFormatTrailingCall:
+        return _calls[_blockCallIndex]._call == child;
+    }
+
+    return true;
+  }
+
+  @override
   void format(CodeWriter writer, State state) {
     switch (state) {
       case State.unsplit:
-        writer.format(_target, allowNewlines: _allowSplitInTarget);
+        writer.format(_target);
 
         for (var i = 0; i < _calls.length; i++) {
           _formatCall(writer, state, i, allowNewlines: false);
@@ -172,7 +193,7 @@
 
       case _splitAfterProperties:
         writer.pushIndent(_indent);
-        writer.format(_target, allowNewlines: _allowSplitInTarget);
+        writer.format(_target);
 
         for (var i = 0; i < _calls.length; i++) {
           writer.splitIf(i >= _leadingProperties, space: false);
@@ -182,7 +203,7 @@
         writer.popIndent();
 
       case _blockFormatTrailingCall:
-        writer.format(_target, allowNewlines: _allowSplitInTarget);
+        writer.format(_target);
 
         for (var i = 0; i < _calls.length; i++) {
           _formatCall(writer, state, i, allowNewlines: i == _blockCallIndex);
@@ -212,8 +233,7 @@
       _ => false,
     };
 
-    writer.format(_calls[i]._call,
-        separate: separate, allowNewlines: allowNewlines);
+    writer.format(_calls[i]._call, separate: separate);
   }
 
   @override
diff --git a/lib/src/piece/clause.dart b/lib/src/piece/clause.dart
index f84902a..33eb18f 100644
--- a/lib/src/piece/clause.dart
+++ b/lib/src/piece/clause.dart
@@ -81,13 +81,25 @@
   final bool _allowLeadingClause;
 
   ClausePiece(this._clauses, {bool allowLeadingClause = false})
-      : _allowLeadingClause = allowLeadingClause;
+      : _allowLeadingClause = allowLeadingClause && _clauses.length > 1;
 
   @override
   List<State> get additionalStates =>
       [if (_allowLeadingClause) _betweenClauses, State.split];
 
   @override
+  bool allowNewlineInChild(State state, Piece child) {
+    if (_allowLeadingClause && child == _clauses.first) {
+      // A split inside the first clause forces a split before the keyword.
+      return state == State.split;
+    } else {
+      // For the other clauses (or if there is no leading one), any split
+      // inside a clause forces all of them to split.
+      return state != State.unsplit;
+    }
+  }
+
+  @override
   void format(CodeWriter writer, State state) {
     writer.pushIndent(Indent.expression);
 
@@ -96,13 +108,13 @@
         // Before the leading clause, only split when in the fully split state.
         // A split inside the first clause forces a split before the keyword.
         writer.splitIf(state == State.split);
-        writer.format(clause, allowNewlines: state == State.split);
+        writer.format(clause);
       } else {
         // For the other clauses (or if there is no leading one), split in the
         // fully split state and any split inside and clause forces all of them
         // to split.
         writer.splitIf(state != State.unsplit);
-        writer.format(clause, allowNewlines: state != State.unsplit);
+        writer.format(clause);
       }
     }
 
diff --git a/lib/src/piece/constructor.dart b/lib/src/piece/constructor.dart
index 58767f3..ed6a070 100644
--- a/lib/src/piece/constructor.dart
+++ b/lib/src/piece/constructor.dart
@@ -127,24 +127,33 @@
   }
 
   @override
-  void format(CodeWriter writer, State state) {
+  bool allowNewlineInChild(State state, Piece child) {
+    if (child == _body) return true;
+
     // 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.
-    var allowNewlines = _initializers == null || state != State.unsplit;
+    return _initializers == null || state != State.unsplit;
+  }
 
-    writer.format(_header, allowNewlines: allowNewlines);
-    writer.format(_parameters, allowNewlines: allowNewlines);
+  @override
+  bool containsNewline(State state) =>
+      state == _splitBeforeInitializers || super.containsNewline(state);
+
+  @override
+  void format(CodeWriter writer, State state) {
+    writer.format(_header);
+    writer.format(_parameters);
 
     if (_redirect case var redirect?) {
       writer.space();
-      writer.format(redirect, allowNewlines: allowNewlines);
+      writer.format(redirect);
     }
 
     if (_initializers case var initializers?) {
       writer.pushIndent(Indent.block);
       writer.splitIf(state == _splitBeforeInitializers);
 
-      writer.format(_initializerSeparator!, allowNewlines: allowNewlines);
+      writer.format(_initializerSeparator!);
       writer.space();
 
       // Indent subsequent initializers past the `:`.
@@ -154,7 +163,7 @@
         writer.pushIndent(Indent.initializer);
       }
 
-      writer.format(initializers, allowNewlines: allowNewlines);
+      writer.format(initializers);
       writer.popIndent();
       writer.popIndent();
     }
diff --git a/lib/src/piece/control_flow.dart b/lib/src/piece/control_flow.dart
index 291fa17..831b420 100644
--- a/lib/src/piece/control_flow.dart
+++ b/lib/src/piece/control_flow.dart
@@ -40,11 +40,17 @@
   }
 
   @override
-  void forEachChild(void Function(Piece piece) callback) {
-    for (var section in _sections) {
-      callback(section.header);
-      callback(section.statement);
+  bool allowNewlineInChild(State state, Piece child) => state == State.split;
+
+  @override
+  bool containsNewline(State state) {
+    if (state == State.split) {
+      for (var section in _sections) {
+        if (!section.isBlock) return true;
+      }
     }
+
+    return super.containsNewline(state);
   }
 
   @override
@@ -53,7 +59,7 @@
       var section = _sections[i];
 
       // A split in the condition forces the branches to split.
-      writer.format(section.header, allowNewlines: state == State.split);
+      writer.format(section.header);
 
       if (!section.isBlock) {
         writer.pushIndent(Indent.block);
@@ -61,7 +67,7 @@
       }
 
       // TODO(perf): Investigate whether it's worth using `separate:` here.
-      writer.format(section.statement, allowNewlines: state == State.split);
+      writer.format(section.statement);
 
       // Reset the indentation for the subsequent `else` or `} else` line.
       if (!section.isBlock) writer.popIndent();
@@ -73,6 +79,14 @@
   }
 
   @override
+  void forEachChild(void Function(Piece piece) callback) {
+    for (var section in _sections) {
+      callback(section.header);
+      callback(section.statement);
+    }
+  }
+
+  @override
   String get debugName => 'Ctrl';
 }
 
diff --git a/lib/src/piece/for.dart b/lib/src/piece/for.dart
index 8144d93..e514606 100644
--- a/lib/src/piece/for.dart
+++ b/lib/src/piece/for.dart
@@ -99,18 +99,22 @@
   List<State> get additionalStates => const [State.split];
 
   @override
+  bool allowNewlineInChild(State state, Piece child) {
+    if (state == State.split) return true;
+
+    // Always allow block-splitting the sequence if it supports it.
+    return child == _sequence && _canBlockSplitSequence;
+  }
+
+  @override
   void format(CodeWriter writer, State state) {
     // When splitting at `in`, both operands may split or not and will be
     // indented if they do.
     if (state == State.split) writer.pushIndent(Indent.expression);
 
-    writer.format(_variable, allowNewlines: state == State.split);
-
+    writer.format(_variable);
     writer.splitIf(state == State.split);
-
-    // Always allow block-splitting the sequence if it supports it.
-    writer.format(_sequence,
-        allowNewlines: _canBlockSplitSequence || state == State.split);
+    writer.format(_sequence);
 
     if (state == State.split) writer.popIndent();
   }
diff --git a/lib/src/piece/if_case.dart b/lib/src/piece/if_case.dart
index 14fcfb1..80d8e08 100644
--- a/lib/src/piece/if_case.dart
+++ b/lib/src/piece/if_case.dart
@@ -70,36 +70,28 @@
       ];
 
   @override
+  bool allowNewlineInChild(State state, Piece child) {
+    return switch (state) {
+      // When not splitting before `case` or `when`, we only allow newlines
+      // in block-formatted patterns.
+      State.unsplit when child == _pattern => _canBlockSplitPattern,
+
+      // Allow newlines only in the guard if we split before `when`.
+      _beforeWhen when child == _guard => true,
+
+      // Only allow the guard on the same line as the pattern if it doesn't
+      // split.
+      _beforeCase when child != _guard => true,
+      _beforeCaseAndWhen => true,
+      _ => false,
+    };
+  }
+
+  @override
   void format(CodeWriter writer, State state) {
-    var allowNewlineInValue = false;
-    var allowNewlineInPattern = false;
-    var allowNewlineInGuard = false;
-
-    switch (state) {
-      case State.unsplit:
-        // When not splitting before `case` or `when`, we only allow newlines
-        // in block-formatted patterns.
-        allowNewlineInPattern = _canBlockSplitPattern;
-
-      case _beforeWhen:
-        // Allow newlines only in the guard if we split before `when`.
-        allowNewlineInGuard = true;
-
-      case _beforeCase:
-        // Only allow the guard on the same line as the pattern if it doesn't
-        // split.
-        allowNewlineInValue = true;
-        allowNewlineInPattern = true;
-
-      case _beforeCaseAndWhen:
-        allowNewlineInValue = true;
-        allowNewlineInPattern = true;
-        allowNewlineInGuard = true;
-    }
-
     if (state != State.unsplit) writer.pushIndent(Indent.expression);
 
-    writer.format(_value, allowNewlines: allowNewlineInValue);
+    writer.format(_value);
 
     // The case clause and pattern.
     writer.splitIf(state == _beforeCase || state == _beforeCaseAndWhen);
@@ -108,14 +100,14 @@
       writer.pushIndent(Indent.expression, canCollapse: true);
     }
 
-    writer.format(_pattern, allowNewlines: allowNewlineInPattern);
+    writer.format(_pattern);
 
     if (!_canBlockSplitPattern) writer.popIndent();
 
     // The guard clause.
     if (_guard case var guard?) {
       writer.splitIf(state == _beforeWhen || state == _beforeCaseAndWhen);
-      writer.format(guard, allowNewlines: allowNewlineInGuard);
+      writer.format(guard);
     }
 
     if (state != State.unsplit) writer.popIndent();
diff --git a/lib/src/piece/infix.dart b/lib/src/piece/infix.dart
index 6672e27..5e7aacc 100644
--- a/lib/src/piece/infix.dart
+++ b/lib/src/piece/infix.dart
@@ -48,10 +48,17 @@
   List<State> get additionalStates => const [State.split];
 
   @override
-  void format(CodeWriter writer, State state) {
+  bool allowNewlineInChild(State state, Piece child) {
+    if (state == State.split) return true;
+
     // Comments before the operands don't force the operator to split.
+    return _leadingComments.contains(child);
+  }
+
+  @override
+  void format(CodeWriter writer, State state) {
     for (var comment in _leadingComments) {
-      writer.format(comment, allowNewlines: true);
+      writer.format(comment);
     }
 
     if (_indent) writer.pushIndent(Indent.expression);
@@ -62,8 +69,7 @@
       // or last operand.
       var separate = state == State.split && i > 0 && i < _operands.length - 1;
 
-      writer.format(_operands[i],
-          separate: separate, allowNewlines: state == State.split);
+      writer.format(_operands[i], separate: separate);
       if (i < _operands.length - 1) writer.splitIf(state == State.split);
     }
 
@@ -82,7 +88,7 @@
 
     for (var operand in _operands) {
       // If any operand contains a newline, then we have to split.
-      if (operand.containsNewline) return State.split;
+      if (operand.containsHardNewline) return State.split;
 
       totalLength += operand.totalCharacters;
       if (totalLength > pageWidth) break;
diff --git a/lib/src/piece/list.dart b/lib/src/piece/list.dart
index 0e8d52a..0ebd7b3 100644
--- a/lib/src/piece/list.dart
+++ b/lib/src/piece/list.dart
@@ -116,6 +116,17 @@
   }
 
   @override
+  bool allowNewlineInChild(State state, Piece child) {
+    if (state == State.split) return true;
+    if (child == _before) return true;
+    if (child == _after) return true;
+
+    // Only some elements (usually a single block element) allow newlines
+    // when the list itself isn't split.
+    return child is ListElementPiece && child.allowNewlinesWhenUnsplit;
+  }
+
+  @override
   void format(CodeWriter writer, State state) {
     // Format the opening bracket, if there is one.
     if (_before case var before?) {
@@ -144,13 +155,7 @@
       var separate = state == State.split &&
           (i > 0 || _before != null) &&
           (i < _elements.length - 1 || _after != null);
-
-      // Only some elements (usually a single block element) allow newlines
-      // when the list itself isn't split.
-      var allowNewlines =
-          element.allowNewlinesWhenUnsplit || state == State.split;
-
-      writer.format(element, separate: separate, allowNewlines: allowNewlines);
+      writer.format(element, separate: separate);
 
       if (state == State.unsplit && element.indentWhenBlockFormatted) {
         writer.popIndent();
@@ -173,7 +178,7 @@
       writer.splitIf(state == State.split,
           space: _style.spaceWhenUnsplit && _elements.isNotEmpty);
 
-      writer.format(after, allowNewlines: true);
+      writer.format(after);
     }
   }
 
@@ -194,7 +199,7 @@
     if (_before case var before?) {
       // A newline in the opening bracket (like a line comment after the
       // bracket) forces the list to split.
-      if (before.containsNewline) return State.split;
+      if (before.containsHardNewline) return State.split;
       totalLength += before.totalCharacters;
     }
 
@@ -203,7 +208,7 @@
       // to split.
       if (element.allowNewlinesWhenUnsplit) continue;
 
-      if (element.containsNewline) return State.split;
+      if (element.containsHardNewline) return State.split;
       totalLength += element.totalCharacters;
       if (totalLength > pageWidth) break;
     }
diff --git a/lib/src/piece/piece.dart b/lib/src/piece/piece.dart
index 174a551..9a45b85 100644
--- a/lib/src/piece/piece.dart
+++ b/lib/src/piece/piece.dart
@@ -43,13 +43,13 @@
   ///
   /// This is lazily computed and cached for performance, so should only be
   /// accessed after all of the piece's children are known.
-  late final bool containsNewline = _calculateContainsNewline();
+  late final bool containsHardNewline = calculateContainsHardNewline();
 
-  bool _calculateContainsNewline() {
+  bool calculateContainsHardNewline() {
     var anyHasNewline = false;
 
     forEachChild((child) {
-      anyHasNewline |= child.containsNewline;
+      anyHasNewline |= child.containsHardNewline;
     });
 
     return anyHasNewline;
@@ -60,9 +60,9 @@
   ///
   /// This is lazily computed and cached for performance, so should only be
   /// accessed after all of the piece's children are known.
-  late final int totalCharacters = _calculateTotalCharacters();
+  late final int totalCharacters = calculateTotalCharacters();
 
-  int _calculateTotalCharacters() {
+  int calculateTotalCharacters() {
     var total = 0;
 
     forEachChild((child) {
@@ -84,6 +84,22 @@
   /// child piece and the state that child should be constrained to.
   void applyConstraints(State state, Constrain constrain) {}
 
+  /// Whether the [child] of this piece should be allowed to contain newlines
+  /// (directly or transitively) when this piece is in [state].
+  bool allowNewlineInChild(State state, Piece child) => true;
+
+  /// Whether this piece contains a newline when this piece is in [state].
+  ///
+  /// This should only return `true` if the piece will *always* write at least
+  /// one newline -- either itself or one of its children -- when in this state.
+  /// If a piece may contain a newline or may not in some state, this should
+  /// return `false`.
+  ///
+  /// By default, we assume that any piece not in [State.unsplit] or that has a
+  /// hard newline will contain a newline.
+  bool containsNewline(State state) =>
+      state != State.unsplit || containsHardNewline;
+
   /// Given that this piece is in [state], use [writer] to produce its formatted
   /// output.
   void format(CodeWriter writer, State state);
@@ -92,8 +108,9 @@
   void forEachChild(void Function(Piece piece) callback);
 
   /// If the piece can determine that it will always end up in a certain state
-  /// given [pageWidth] and size metrics returned by calling [containsNewline]
-  /// and [totalCharacters] on its children, then returns that [State].
+  /// given [pageWidth] and size metrics returned by calling
+  /// [containsHardNewline] and [totalCharacters] on its children, then returns
+  /// that [State].
   ///
   /// For example, a series of infix operators wider than a page will always
   /// split one per operator. If we can determine this eagerly just based on
@@ -146,6 +163,29 @@
     pin(State.unsplit);
   }
 
+  /// All of the transitive children of this piece (including the piece itself)
+  /// that have more than state.
+  ///
+  /// This calculated and cached because it's faster than traversing the child
+  /// tree and having to skip past all of the stateless [AdjacentPiece],
+  /// [SpacePiece], [SequencePiece], etc.
+  late final List<Piece> statefulOffspring = _calculateStatelessOffspring();
+
+  List<Piece> _calculateStatelessOffspring() {
+    // TODO(rnystrom): Traversing and storing the transitive list of stateless
+    // offspring at every level of the Piece tree might be slow. Is it? If so,
+    // is there a faster way to propagate constraints to the relevant parts of
+    // the subtree?
+    var result = <Piece>[];
+    void traverse(Piece piece) {
+      if (piece.additionalStates.isNotEmpty) result.add(piece);
+      piece.forEachChild(traverse);
+    }
+
+    traverse(this);
+    return result;
+  }
+
   /// The name of this piece as it appears in debug output.
   ///
   /// By default, this is the class's name with `Piece` removed.
@@ -155,198 +195,6 @@
   String toString() => '$debugName${_pinnedState ?? ''}';
 }
 
-/// A simple atomic piece of code.
-///
-/// This may represent a series of tokens where no split can occur between them.
-/// It may also contain one or more comments.
-sealed class TextPiece extends Piece {
-  /// RegExp that matches any valid Dart line terminator.
-  static final _lineTerminatorPattern = RegExp(r'\r\n?|\n');
-
-  /// The lines of text in this piece.
-  ///
-  /// Most [TextPieces] will contain only a single line, but a piece for a
-  /// multi-line string or comment will have multiple lines. These are stored
-  /// as separate lines instead of a single multi-line Dart String so that
-  /// line endings are normalized and so that column calculation during line
-  /// splitting calculates each line in the piece separately.
-  final List<String> _lines = [''];
-
-  /// The offset from the beginning of [text] where the selection starts, or
-  /// `null` if the selection does not start within this chunk.
-  int? _selectionStart;
-
-  /// The offset from the beginning of [text] where the selection ends, or
-  /// `null` if the selection does not start within this chunk.
-  int? _selectionEnd;
-
-  /// Append [text] to the end of this piece.
-  ///
-  /// If [text] may contain any newline characters, then [multiline] must be
-  /// `true`.
-  ///
-  /// If [selectionStart] and/or [selectionEnd] are given, then notes that the
-  /// corresponding selection markers appear that many code units from where
-  /// [text] will be appended.
-  void append(String text,
-      {bool multiline = false, int? selectionStart, int? selectionEnd}) {
-    if (selectionStart != null) {
-      _selectionStart = _adjustSelection(selectionStart);
-    }
-
-    if (selectionEnd != null) {
-      _selectionEnd = _adjustSelection(selectionEnd);
-    }
-
-    if (multiline) {
-      var lines = text.split(_lineTerminatorPattern);
-      for (var i = 0; i < lines.length; i++) {
-        if (i > 0) _lines.add('');
-        _lines.last += lines[i];
-      }
-    } else {
-      _lines.last += text;
-    }
-  }
-
-  /// Sets [selectionStart] to be [start] code units after the end of the
-  /// current text in this piece.
-  void startSelection(int start) {
-    _selectionStart = _adjustSelection(start);
-  }
-
-  /// Sets [selectionEnd] to be [end] code units after the end of the
-  /// current text in this piece.
-  void endSelection(int end) {
-    _selectionEnd = _adjustSelection(end);
-  }
-
-  /// Adjust [offset] by the current length of this [TextPiece].
-  int _adjustSelection(int offset) {
-    for (var line in _lines) {
-      offset += line.length;
-    }
-
-    return offset;
-  }
-
-  void _formatSelection(CodeWriter writer) {
-    if (_selectionStart case var start?) {
-      writer.startSelection(start);
-    }
-
-    if (_selectionEnd case var end?) {
-      writer.endSelection(end);
-    }
-  }
-
-  void _formatLines(CodeWriter writer) {
-    for (var i = 0; i < _lines.length; i++) {
-      if (i > 0) writer.newline(flushLeft: i > 0);
-      writer.write(_lines[i]);
-    }
-  }
-
-  @override
-  bool _calculateContainsNewline() => _lines.length > 1;
-
-  @override
-  int _calculateTotalCharacters() {
-    var total = 0;
-
-    for (var line in _lines) {
-      total += line.length;
-    }
-
-    return total;
-  }
-
-  @override
-  String toString() => '`${_lines.join('¬')}`';
-}
-
-/// [TextPiece] for non-comment source code that may have comments attached to
-/// it.
-class CodePiece extends TextPiece {
-  /// Pieces for any comments that appear immediately before this code.
-  final List<Piece> _leadingComments;
-
-  /// Pieces for any comments that hang off the same line as this code.
-  final List<Piece> _hangingComments = [];
-
-  CodePiece([this._leadingComments = const []]);
-
-  void addHangingComment(Piece comment) {
-    _hangingComments.add(comment);
-  }
-
-  @override
-  void format(CodeWriter writer, State state) {
-    _formatSelection(writer);
-
-    if (_leadingComments.isNotEmpty) {
-      // Always put leading comments on a new line.
-      writer.newline();
-
-      for (var comment in _leadingComments) {
-        writer.format(comment);
-      }
-    }
-
-    _formatLines(writer);
-
-    for (var comment in _hangingComments) {
-      writer.space();
-      writer.format(comment);
-    }
-  }
-
-  @override
-  void forEachChild(void Function(Piece piece) callback) {
-    _leadingComments.forEach(callback);
-    _hangingComments.forEach(callback);
-  }
-}
-
-/// A [TextPiece] for a source code comment and the whitespace after it, if any.
-class CommentPiece extends TextPiece {
-  /// Whitespace at the end of the comment.
-  final Whitespace _trailingWhitespace;
-
-  CommentPiece([this._trailingWhitespace = Whitespace.none]);
-
-  @override
-  void format(CodeWriter writer, State state) {
-    _formatSelection(writer);
-    _formatLines(writer);
-    writer.whitespace(_trailingWhitespace);
-  }
-
-  @override
-  bool _calculateContainsNewline() =>
-      _trailingWhitespace.hasNewline || super._calculateContainsNewline();
-
-  @override
-  void forEachChild(void Function(Piece piece) callback) {}
-}
-
-/// A piece that writes a single space.
-class SpacePiece extends Piece {
-  @override
-  void forEachChild(void Function(Piece piece) callback) {}
-
-  @override
-  void format(CodeWriter writer, State state) {
-    writer.space();
-  }
-
-  @override
-  bool _calculateContainsNewline() => false;
-
-  @override
-  int _calculateTotalCharacters() => 1;
-}
-
 /// A state that a piece can be in.
 ///
 /// Each state identifies one way that a piece can be split into multiple lines.
diff --git a/lib/src/piece/sequence.dart b/lib/src/piece/sequence.dart
index 71fc8a5..2df4911 100644
--- a/lib/src/piece/sequence.dart
+++ b/lib/src/piece/sequence.dart
@@ -3,80 +3,97 @@
 // BSD-style license that can be found in the LICENSE file.
 
 import '../back_end/code_writer.dart';
+import '../constants.dart';
 import 'piece.dart';
 
 /// A piece for a series of statements or members inside a block or declaration
-/// body.
+/// body or at the top level of a program.
 ///
-/// Usually constructed using a [SequenceBuilder].
+/// Constructed using a [SequenceBuilder].
 class SequencePiece extends Piece {
-  /// The opening delimiter, if any.
-  final Piece? _leftBracket;
-
   /// The series of members or statements.
   final List<SequenceElementPiece> _elements;
 
-  SequencePiece(this._elements, {Piece? leftBracket, Piece? rightBracket})
-      : _leftBracket = leftBracket,
-        _rightBracket = rightBracket;
-
-  /// The closing delimiter, if any.
-  final Piece? _rightBracket;
-
-  @override
-  List<State> get additionalStates => [if (_elements.isNotEmpty) State.split];
+  SequencePiece(this._elements);
 
   @override
   void format(CodeWriter writer, State state) {
-    if (_leftBracket case var leftBracket?) {
-      writer.format(leftBracket, allowNewlines: state == State.split);
-      writer.pushIndent(_elements.firstOrNull?._indent ?? 0);
-      writer.splitIf(state == State.split, space: false);
-    }
+    writer.pushIndent(Indent.none);
 
     for (var i = 0; i < _elements.length; i++) {
       var element = _elements[i];
 
-      // We can format an element separately if the element is on its own line.
-      // This happens when the sequence is split and there is something before
-      // and after the element, either brackets or other items.
-      var separate = state == State.split &&
-          (i > 0 || _leftBracket != null) &&
-          (i < _elements.length - 1 || _rightBracket != null);
-
-      writer.format(element,
-          separate: separate, allowNewlines: state == State.split);
+      writer.format(element, separate: true);
 
       if (i < _elements.length - 1) {
-        if (_leftBracket != null || i > 0) writer.popIndent();
+        writer.popIndent();
         writer.pushIndent(_elements[i + 1]._indent);
         writer.newline(blank: element.blankAfter);
       }
     }
 
-    if (_leftBracket != null || _elements.length > 1) writer.popIndent();
-
-    if (_rightBracket case var rightBracket?) {
-      writer.splitIf(state == State.split, space: false);
-      writer.format(rightBracket, allowNewlines: state == State.split);
-    }
+    writer.popIndent();
   }
 
   @override
   void forEachChild(void Function(Piece piece) callback) {
-    if (_leftBracket case var leftBracket?) callback(leftBracket);
-
     for (var element in _elements) {
       callback(element);
     }
-
-    if (_rightBracket case var rightBracket?) callback(rightBracket);
   }
 
+  /// If there are multiple elements, there are newlines between them.
+  @override
+  bool calculateContainsHardNewline() => _elements.length > 1;
+
   @override
   String get debugName => 'Seq';
 }
 
+/// A piece for a non-empty brace-delimited series of statements or members
+/// inside a block or declaration body.
+///
+/// Unlike [ListPiece], always splits between the elements.
+///
+/// Constructed using a [SequenceBuilder].
+class BlockPiece extends Piece {
+  /// The opening delimiter.
+  final Piece _leftBracket;
+
+  /// The series of members or statements.
+  final SequencePiece _elements;
+
+  /// The closing delimiter.
+  final Piece _rightBracket;
+
+  BlockPiece(this._leftBracket, this._elements, this._rightBracket);
+
+  @override
+  void format(CodeWriter writer, State state) {
+    writer.format(_leftBracket);
+    writer.pushIndent(Indent.block);
+    writer.newline();
+    writer.format(_elements);
+    writer.popIndent();
+    writer.newline();
+    writer.format(_rightBracket);
+  }
+
+  @override
+  void forEachChild(void Function(Piece piece) callback) {
+    callback(_leftBracket);
+    callback(_elements);
+    callback(_rightBracket);
+  }
+
+  /// A [BlockPiece] is never empty and always splits between the delimiters.
+  @override
+  bool calculateContainsHardNewline() => true;
+
+  @override
+  String get debugName => 'Block';
+}
+
 /// An element inside a [SequencePiece].
 ///
 /// Tracks the underlying [Piece] along with surrounding whitespace.
diff --git a/lib/src/piece/text.dart b/lib/src/piece/text.dart
new file mode 100644
index 0000000..274516c
--- /dev/null
+++ b/lib/src/piece/text.dart
@@ -0,0 +1,215 @@
+// 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.
+
+import '../back_end/code_writer.dart';
+import 'piece.dart';
+
+/// A simple atomic piece of code.
+///
+/// This may represent a series of tokens where no split can occur between them.
+/// It may also contain one or more comments.
+sealed class TextPiece extends Piece {
+  /// RegExp that matches any valid Dart line terminator.
+  static final _lineTerminatorPattern = RegExp(r'\r\n?|\n');
+
+  /// The lines of text in this piece.
+  ///
+  /// Most [TextPieces] will contain only a single line, but a piece for a
+  /// multi-line string or comment will have multiple lines. These are stored
+  /// as separate lines instead of a single multi-line Dart String so that
+  /// line endings are normalized and so that column calculation during line
+  /// splitting calculates each line in the piece separately.
+  final List<String> _lines = [''];
+
+  /// The offset from the beginning of [text] where the selection starts, or
+  /// `null` if the selection does not start within this chunk.
+  int? _selectionStart;
+
+  /// The offset from the beginning of [text] where the selection ends, or
+  /// `null` if the selection does not start within this chunk.
+  int? _selectionEnd;
+
+  /// Append [text] to the end of this piece.
+  ///
+  /// If [text] may contain any newline characters, then [multiline] must be
+  /// `true`.
+  ///
+  /// If [selectionStart] and/or [selectionEnd] are given, then notes that the
+  /// corresponding selection markers appear that many code units from where
+  /// [text] will be appended.
+  void append(String text,
+      {bool multiline = false, int? selectionStart, int? selectionEnd}) {
+    if (selectionStart != null) {
+      _selectionStart = _adjustSelection(selectionStart);
+    }
+
+    if (selectionEnd != null) {
+      _selectionEnd = _adjustSelection(selectionEnd);
+    }
+
+    if (multiline) {
+      var lines = text.split(_lineTerminatorPattern);
+      for (var i = 0; i < lines.length; i++) {
+        if (i > 0) _lines.add('');
+        _lines.last += lines[i];
+      }
+    } else {
+      _lines.last += text;
+    }
+  }
+
+  /// Sets [selectionStart] to be [start] code units after the end of the
+  /// current text in this piece.
+  void startSelection(int start) {
+    _selectionStart = _adjustSelection(start);
+  }
+
+  /// Sets [selectionEnd] to be [end] code units after the end of the
+  /// current text in this piece.
+  void endSelection(int end) {
+    _selectionEnd = _adjustSelection(end);
+  }
+
+  /// Adjust [offset] by the current length of this [TextPiece].
+  int _adjustSelection(int offset) {
+    for (var line in _lines) {
+      offset += line.length;
+    }
+
+    return offset;
+  }
+
+  void _formatSelection(CodeWriter writer) {
+    if (_selectionStart case var start?) {
+      writer.startSelection(start);
+    }
+
+    if (_selectionEnd case var end?) {
+      writer.endSelection(end);
+    }
+  }
+
+  void _formatLines(CodeWriter writer) {
+    for (var i = 0; i < _lines.length; i++) {
+      if (i > 0) writer.newline(flushLeft: i > 0);
+      writer.write(_lines[i]);
+    }
+  }
+
+  @override
+  bool calculateContainsHardNewline() => _lines.length > 1;
+
+  @override
+  int calculateTotalCharacters() {
+    var total = 0;
+
+    for (var line in _lines) {
+      total += line.length;
+    }
+
+    return total;
+  }
+
+  @override
+  String toString() => '`${_lines.join('¬')}`';
+}
+
+/// [TextPiece] for non-comment source code that may have comments attached to
+/// it.
+class CodePiece extends TextPiece {
+  /// Pieces for any comments that appear immediately before this code.
+  final List<Piece> _leadingComments;
+
+  /// Pieces for any comments that hang off the same line as this code.
+  final List<Piece> _hangingComments = [];
+
+  CodePiece([this._leadingComments = const []]);
+
+  void addHangingComment(Piece comment) {
+    _hangingComments.add(comment);
+  }
+
+  @override
+  void format(CodeWriter writer, State state) {
+    _formatSelection(writer);
+
+    if (_leadingComments.isNotEmpty) {
+      // Always put leading comments on a new line.
+      writer.newline();
+
+      for (var comment in _leadingComments) {
+        writer.format(comment);
+      }
+    }
+
+    _formatLines(writer);
+
+    for (var comment in _hangingComments) {
+      writer.space();
+      writer.format(comment);
+    }
+  }
+
+  @override
+  void forEachChild(void Function(Piece piece) callback) {
+    _leadingComments.forEach(callback);
+    _hangingComments.forEach(callback);
+  }
+}
+
+/// A [TextPiece] for a source code comment and the whitespace after it, if any.
+class CommentPiece extends TextPiece {
+  /// Whitespace at the end of the comment.
+  final Whitespace _trailingWhitespace;
+
+  CommentPiece([this._trailingWhitespace = Whitespace.none]);
+
+  @override
+  void format(CodeWriter writer, State state) {
+    _formatSelection(writer);
+    _formatLines(writer);
+    writer.whitespace(_trailingWhitespace);
+  }
+
+  @override
+  bool calculateContainsHardNewline() =>
+      _trailingWhitespace.hasNewline || super.calculateContainsHardNewline();
+
+  @override
+  void forEachChild(void Function(Piece piece) callback) {}
+}
+
+/// A piece that writes a single space.
+class SpacePiece extends Piece {
+  @override
+  void forEachChild(void Function(Piece piece) callback) {}
+
+  @override
+  void format(CodeWriter writer, State state) {
+    writer.space();
+  }
+
+  @override
+  bool calculateContainsHardNewline() => false;
+
+  @override
+  int calculateTotalCharacters() => 1;
+}
+
+/// A piece that writes a single newline.
+class NewlinePiece extends Piece {
+  @override
+  void forEachChild(void Function(Piece piece) callback) {}
+
+  @override
+  void format(CodeWriter writer, State state) {
+    writer.newline();
+  }
+
+  @override
+  bool calculateContainsHardNewline() => true;
+
+  @override
+  int calculateTotalCharacters() => 0;
+}
diff --git a/lib/src/piece/type.dart b/lib/src/piece/type.dart
index df1259c..3e2cb6c 100644
--- a/lib/src/piece/type.dart
+++ b/lib/src/piece/type.dart
@@ -36,15 +36,19 @@
   }
 
   @override
-  void format(CodeWriter writer, State state) {
+  bool allowNewlineInChild(State state, Piece child) {
+    if (child == _body) return true;
+
     // If the body may or may not split, then a newline in the header or
     // clauses forces the body to split.
-    var allowSplitInHeader =
-        _bodyType != TypeBodyType.list || state == State.split;
+    return _bodyType != TypeBodyType.list || state == State.split;
+  }
 
-    writer.format(_header, allowNewlines: allowSplitInHeader);
+  @override
+  void format(CodeWriter writer, State state) {
+    writer.format(_header);
     if (_clauses case var clauses?) {
-      writer.format(clauses, allowNewlines: allowSplitInHeader);
+      writer.format(clauses);
     }
 
     if (_bodyType != TypeBodyType.semicolon) writer.space();
diff --git a/lib/src/piece/variable.dart b/lib/src/piece/variable.dart
index 90738ec..97dfe29 100644
--- a/lib/src/piece/variable.dart
+++ b/lib/src/piece/variable.dart
@@ -60,8 +60,17 @@
       ];
 
   @override
+  bool allowNewlineInChild(State state, Piece child) {
+    if (child == _header) {
+      return state != State.unsplit;
+    } else {
+      return _variables.length == 1 || state != State.unsplit;
+    }
+  }
+
+  @override
   void format(CodeWriter writer, State state) {
-    writer.format(_header, allowNewlines: state != State.unsplit);
+    writer.format(_header);
 
     // If we split at the variables (but not the type), then indent the
     // variables and their initializers.
@@ -75,9 +84,7 @@
       if (i > 0) writer.splitIf(state != State.unsplit);
 
       // TODO(perf): Investigate whether it's worth using `separate:` here.
-      // Force multiple variables to split if an initializer does.
-      writer.format(_variables[i],
-          allowNewlines: _variables.length == 1 || state != State.unsplit);
+      writer.format(_variables[i]);
     }
 
     if (state == _betweenVariables) writer.popIndent();
diff --git a/test/tall/declaration/class_comment.unit b/test/tall/declaration/class_comment.unit
index 436141e..43a96b2 100644
--- a/test/tall/declaration/class_comment.unit
+++ b/test/tall/declaration/class_comment.unit
@@ -20,7 +20,9 @@
 >>> Empty class containing inline block comment.
 class C {   /* comment */  }
 <<<
-class C {/* comment */}
+class C {
+  /* comment */
+}
 >>> Empty class containing non-inline block comment.
 class C {
 
diff --git a/test/tall/declaration/extension_type_comment.unit b/test/tall/declaration/extension_type_comment.unit
index 8fbce61..5ece81c 100644
--- a/test/tall/declaration/extension_type_comment.unit
+++ b/test/tall/declaration/extension_type_comment.unit
@@ -12,7 +12,9 @@
 ) /*j*/
     implements /*k*/
         I1 /*l*/, /*m*/
-        I2 /*n*/ {/*o*/} /*p*/
+        I2 /*n*/ {
+  /*o*/
+} /*p*/
 >>> Line comments everywhere.
 // 0
 @patch // a
diff --git a/test/tall/statement/block_comment.stmt b/test/tall/statement/block_comment.stmt
index c8cbfe2..eaf2e0b 100644
--- a/test/tall/statement/block_comment.stmt
+++ b/test/tall/statement/block_comment.stmt
@@ -9,7 +9,9 @@
 >>> Empty block containing inline block comment.
 {   /* comment */  }
 <<<
-{/* comment */}
+{
+  /* comment */
+}
 >>> Empty block containing non-inline block comment.
 {
 
diff --git a/test/tall/statement/switch_comment.stmt b/test/tall/statement/switch_comment.stmt
index ce53767..26ee9fc 100644
--- a/test/tall/statement/switch_comment.stmt
+++ b/test/tall/statement/switch_comment.stmt
@@ -135,13 +135,19 @@
 switch (e) {/* comment */
 }
 <<<
-switch (e) {/* comment */}
+switch (e) {
+  /* comment */
+}
 >>> Block comment with leading newline.
 switch (e) {
   /* comment */}
 <<<
-switch (e) {/* comment */}
+switch (e) {
+  /* comment */
+}
 >>> Inline block comment.
 switch (e) {  /* comment */  }
 <<<
-switch (e) {/* comment */}
\ No newline at end of file
+switch (e) {
+  /* comment */
+}
\ No newline at end of file