Don't use AssignPiece for for-in loop "in" clauses. (#1485)

Don't use AssignPiece for for-in loop "in" clauses.

One of the reasons that AssignPiece is so hairy is because it handles
both splitting after the operator ("=", ":", etc.) and before "in".

Splitting out for-in loops to their own Piece class lets us get rid of
a little of that complexity.

Also, I thought that was the only reason that the operator is a separate
Piece. That's actually not true because we also need the LHS to be its
own piece for applyConstraints().

Even so, I still think it's a little easier to understand with for-in
loops being their own piece.

I also slightly tweaked the formatting of for-in loops. In the rare
case where the left side block splits, we now indent it and split at
"in". This is consistent with the old style. I don't have a strong
preference, but I think it looks a little neater this way.

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

Author: munificent

lib/src/front_end/ast_node_visitor.dart

@@ -1204,8 +1204,7 @@ class AstNodeVisitor extends ThrowingAstVisitor<void> with PieceFactory {
 
   @override
   void visitMapLiteralEntry(MapLiteralEntry node) {
-    writeAssignment(node.key, node.separator, node.value,
-        spaceBeforeOperator: false);
+    writeAssignment(node.key, node.separator, node.value);
   }
 
   @override
@@ -1220,8 +1219,7 @@ class AstNodeVisitor extends ThrowingAstVisitor<void> with PieceFactory {
 
   @override
   void visitMapPatternEntry(MapPatternEntry node) {
-    writeAssignment(node.key, node.separator, node.value,
-        spaceBeforeOperator: false);
+    writeAssignment(node.key, node.separator, node.value);
   }
 
   @override
@@ -1271,8 +1269,7 @@ class AstNodeVisitor extends ThrowingAstVisitor<void> with PieceFactory {
 
   @override
   void visitNamedExpression(NamedExpression node) {
-    writeAssignment(node.name.label, node.name.colon, node.expression,
-        spaceBeforeOperator: false);
+    writeAssignment(node.name.label, node.name.colon, node.expression);
   }
 
   @override

lib/src/front_end/piece_factory.dart

@@ -473,10 +473,7 @@ mixin PieceFactory {
         // statement or element.
         forPartsPiece = pieces.build(() {
           pieces.token(leftParenthesis);
-          writeAssignment(
-              variable, forEachParts.inKeyword, forEachParts.iterable,
-              splitBeforeOperator: true,
-              leftHandSideContext: NodeContext.forLoopVariable);
+          writeForIn(variable, forEachParts.inKeyword, forEachParts.iterable);
           pieces.token(rightParenthesis);
         });
 
@@ -491,10 +488,7 @@ mixin PieceFactory {
             pieces.token(keyword);
             pieces.space();
 
-            writeAssignment(
-                pattern, forEachParts.inKeyword, forEachParts.iterable,
-                splitBeforeOperator: true,
-                leftHandSideContext: NodeContext.forLoopVariable);
+            writeForIn(pattern, forEachParts.inKeyword, forEachParts.iterable);
           });
           pieces.token(rightParenthesis);
         });
@@ -1299,11 +1293,6 @@ mixin PieceFactory {
   /// * Expression (`=>`) function body
   /// * Named argument or named record field (`:`)
   /// * Map entry (`:`)
-  /// * For-in loop iterator (`in`)
-  ///
-  /// If [splitBeforeOperator] is `true`, then puts [operator] at the beginning
-  /// of the next line when it splits. Otherwise, puts the operator at the end
-  /// of the preceding line.
   ///
   /// If [canBlockSplitLeft] is `true`, then the left-hand operand supports
   /// being block-formatted without indenting it farther, like:
@@ -1313,9 +1302,7 @@ mixin PieceFactory {
   ///     ] = list;
   void writeAssignment(
       AstNode leftHandSide, Token operator, AstNode rightHandSide,
-      {bool splitBeforeOperator = false,
-      bool includeComma = false,
-      bool spaceBeforeOperator = true,
+      {bool includeComma = false,
       bool canBlockSplitLeft = false,
       NodeContext leftHandSideContext = NodeContext.none}) {
     // If an operand can have block formatting, then a newline in it doesn't
@@ -1345,25 +1332,38 @@ mixin PieceFactory {
     var leftPiece = nodePiece(leftHandSide, context: leftHandSideContext);
 
     var operatorPiece = pieces.build(() {
-      if (spaceBeforeOperator) pieces.space();
+      if (operator.type != TokenType.COLON) pieces.space();
       pieces.token(operator);
-      if (splitBeforeOperator) pieces.space();
     });
 
     var rightPiece = nodePiece(rightHandSide,
-        commaAfter: includeComma,
-        context:
-            splitBeforeOperator ? NodeContext.none : NodeContext.assignment);
+        commaAfter: includeComma, context: NodeContext.assignment);
 
     pieces.add(AssignPiece(
         left: leftPiece,
         operatorPiece,
         rightPiece,
-        splitBeforeOperator: splitBeforeOperator,
         canBlockSplitLeft: canBlockSplitLeft,
         canBlockSplitRight: canBlockSplitRight));
   }
 
+  /// Writes a [Piece] for the `<variable> in <expression>` part of a for-in
+  /// loop.
+  void writeForIn(AstNode leftHandSide, Token inKeyword, Expression sequence) {
+    var leftPiece =
+        nodePiece(leftHandSide, context: NodeContext.forLoopVariable);
+
+    var sequencePiece = pieces.build(() {
+      // Put the `in` at the beginning of the sequence.
+      pieces.token(inKeyword);
+      pieces.space();
+      pieces.visit(sequence);
+    });
+
+    pieces.add(ForInPiece(leftPiece, sequencePiece,
+        canBlockSplitSequence: sequence.canBlockSplit));
+  }
+
   /// Writes a piece for a parameter-like constructor: Either a simple formal
   /// parameter or a record type field, which is syntactically similar to a
   /// parameter.

lib/src/piece/assign.dart

@@ -5,13 +5,15 @@ import '../back_end/code_writer.dart';
 import '../constants.dart';
 import 'piece.dart';
 
-/// A piece for any construct where `=` is followed by an expression: variable
-/// initializer, assignment, constructor initializer, etc.
+/// A piece for an assignment-like construct where an operator is followed by
+/// an expression but where the left side of the operator isn't also an
+/// expression. Used for:
 ///
-/// This piece is also used for map entries and named arguments where `:` is
-/// followed by an expression or element because those also want to support the
-/// "block-like" formatting of delimited expressions on the right, and for the
-/// `in` clause in for-in loops.
+/// - Assignment (`=`, `+=`, etc.)
+/// - Named arguments (`:`)
+/// - Map entries (`:`)
+/// - Record fields (`:`)
+/// - Expression function bodies (`=>`)
 ///
 /// These constructs can be formatted four ways:
 ///
@@ -76,25 +78,21 @@ class AssignPiece extends Piece {
   /// Split at the operator.
   static const State _atOperator = State(3);
 
-  /// The left-hand side of the operation. Includes the operator unless it is
-  /// `in`.
+  /// The left-hand side of the operation.
   final Piece? _left;
 
-  // TODO(perf): Most AssignPieces don't allow splitting between [_left] and
-  // [_operator]. Also, in the common case where the AssignPiece is for a named
-  // argument, then both [_left] and [_operator] will be simple CodePieces. If
-  // we used a single piece for both, they can often be concatenated into a
-  // single [CodePiece]. We only store [_operator] separately for for-in loops.
-  // Consider handling those with a separate Piece class and merging [_left]
-  // and [_operator] in this one.
+  // TODO(rnystrom): If it wasn't for the need to constrain [_left] to split
+  // in [applyConstraints()], we could write the operator into the same piece
+  // as [_left]. In the common case where the AssignPiece is for a named
+  // argument, the name and `:` would then end up in a single atomic
+  // [CodePiece].
 
+  /// The `=` or other operator.
   final Piece _operator;
 
   /// The right-hand side of the operation.
   final Piece _right;
 
-  final bool _splitBeforeOperator;
-
   /// If `true`, then the left side supports being block-formatted, like:
   ///
   ///     var [
@@ -113,11 +111,9 @@ class AssignPiece extends Piece {
 
   AssignPiece(this._operator, this._right,
       {Piece? left,
-      bool splitBeforeOperator = false,
       bool canBlockSplitLeft = false,
       bool canBlockSplitRight = false})
       : _left = left,
-        _splitBeforeOperator = splitBeforeOperator,
         _canBlockSplitLeft = canBlockSplitLeft,
         _canBlockSplitRight = canBlockSplitRight;
 
@@ -130,10 +126,16 @@ class AssignPiece extends Piece {
         _atOperator,
       ];
 
-  /// Apply constraints between how the parameters may split and how the
-  /// initializers may split.
   @override
   void applyConstraints(State state, Constrain constrain) {
+    // Force the left side to block split when in that state.
+    //
+    // Otherwise, the solver may instead leave it unsplit and then split the
+    // right side incorrectly as in:
+    //
+    //  (x, y) = longOperand2 +
+    //      longOperand2 +
+    //      longOperand3;
     if (state == _blockSplitLeft) constrain(_left!, State.split);
   }
 
@@ -172,13 +174,10 @@ class AssignPiece extends Piece {
   }
 
   void _writeOperator(CodeWriter writer, {bool split = false}) {
-    if (_splitBeforeOperator) writer.splitIf(split);
-
     writer.pushIndent(Indent.expression);
     writer.format(_operator);
     writer.popIndent();
-
-    if (!_splitBeforeOperator) writer.splitIf(split);
+    writer.splitIf(split);
   }
 
   void _writeRight(CodeWriter writer,

lib/src/piece/for.dart

@@ -80,3 +80,78 @@ class ForPiece extends Piece {
     callback(_body);
   }
 }
+
+/// A piece for the `<variable> in <expression>` part of a for-in loop.
+///
+/// Can be formatted two ways:
+///
+/// [State.unsplit] No split at all:
+///
+///     for (var x in y) ...
+///
+/// This state also allows splitting the sequence expression if it can be block
+/// formatted:
+///
+///     for (var i in [
+///       element1,
+///       element2,
+///       element3,
+///     ];
+///
+/// [State.split] Split at the `in` operator and allow expression splitting on
+/// either side. Allows:
+///
+///     for (var (longVariable &&
+///             anotherVariable)
+///         in longOperand +
+///             anotherOperand) {
+///       ...
+///     }
+class ForInPiece extends Piece {
+  /// The variable or pattern initialized with each loop iteration.
+  final Piece _variable;
+
+  /// The `in` keyword followed by the sequence expression.
+  final Piece _sequence;
+
+  /// If `true` then the sequence expression supports being block-formatted,
+  /// like:
+  ///
+  ///     for (var e in [
+  ///       element1,
+  ///       element2,
+  ///     ]) {
+  ///       // ...
+  ///     }
+  final bool _canBlockSplitSequence;
+
+  ForInPiece(this._variable, this._sequence,
+      {bool canBlockSplitSequence = false})
+      : _canBlockSplitSequence = canBlockSplitSequence;
+
+  @override
+  List<State> get additionalStates => const [State.split];
+
+  @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.splitIf(state == State.split);
+
+    // Always allow block-splitting the sequence if it supports it.
+    writer.format(_sequence,
+        allowNewlines: _canBlockSplitSequence || state == State.split);
+
+    if (state == State.split) writer.popIndent();
+  }
+
+  @override
+  void forEachChild(void Function(Piece piece) callback) {
+    callback(_variable);
+    callback(_sequence);
+  }
+}

test/tall/statement/for_in.stmt

@@ -113,9 +113,10 @@ for (var (longIdentifier &&
 for (var [longIdentifier, anotherLongOne] in obj) {;}
 <<<
 for (var [
-  longIdentifier,
-  anotherLongOne,
-] in obj) {
+      longIdentifier,
+      anotherLongOne,
+    ]
+    in obj) {
   ;
 }
 >>> With pattern, split in value.