Format if elements. (#1344)

Format if elements.

I was hoping to be able to reuse more of the createIf() code in
PieceFactory for both if elements and if statements, but there are
enough differences between the two that that didn't work out, so I just
inlined the code for both of them directly in AstNodeVisitor.

I was able to reuse the existing IfPiece, though, which is nice.

Author: munificent

lib/src/ast_extensions.dart

@@ -93,6 +93,23 @@ extension AstNodeExtensions on AstNode {
 
     return null;
   }
+
+  /// If this is a spread of a non-empty collection literal, then returns `this`
+  /// as a [SpreadElement].
+  ///
+  /// Otherwise, returns `null`.
+  SpreadElement? get spreadCollection {
+    var node = this;
+    if (node is! SpreadElement) return null;
+
+    return switch (node.expression) {
+      ListLiteral(:var elements, :var rightBracket) ||
+      SetOrMapLiteral(:var elements, :var rightBracket)
+          when elements.canSplit(rightBracket) =>
+        node,
+      _ => null,
+    };
+  }
 }
 
 extension AstIterableExtensions on Iterable<AstNode> {

lib/src/back_end/solution.dart

@@ -160,7 +160,7 @@ class Solution implements Comparable<Solution> {
     if (_nextPieceToExpand case var piece?) {
       return [
         for (var state in piece.states)
-          if (_stateSet.tryBind(piece, state) case final stateSet?)
+          if (_stateSet.tryBind(piece, state) case var stateSet?)
             Solution._(root, pageWidth, stateSet)
       ];
     }

lib/src/front_end/ast_node_visitor.dart

@@ -935,12 +935,151 @@ class AstNodeVisitor extends ThrowingAstVisitor<Piece> with PieceFactory {
 
   @override
   Piece visitIfElement(IfElement node) {
-    throw UnimplementedError();
+    var piece = IfPiece();
+
+    // Recurses through the else branches to flatten them into a linear if-else
+    // chain handled by a single [IfPiece].
+    void traverse(Token? precedingElse, IfElement ifElement) {
+      var spreadThen = ifElement.thenElement.spreadCollection;
+
+      var condition = buildPiece((b) {
+        b.token(precedingElse, spaceAfter: true);
+        b.add(startControlFlow(ifElement.ifKeyword, ifElement.leftParenthesis,
+            ifElement.expression, ifElement.rightParenthesis));
+
+        // Make the `...` part of the header so that IfPiece can correctly
+        // constrain the inner collection literal's ListPiece to split.
+        if (spreadThen != null) {
+          b.space();
+          b.token(spreadThen.spreadOperator);
+        }
+      });
+
+      Piece thenElement;
+      if (spreadThen != null) {
+        thenElement = nodePiece(spreadThen.expression);
+      } else {
+        thenElement = nodePiece(ifElement.thenElement);
+      }
+
+      // If the then branch of an if element is itself a control flow
+      // element, then force the outer if to always split.
+      if (ifElement.thenElement.isControlFlowElement) {
+        piece.pin(State.split);
+      }
+
+      piece.add(condition, thenElement, isBlock: spreadThen != null);
+
+      switch (ifElement.elseElement) {
+        case IfElement elseIf:
+          // Hit an else-if, so flatten it into the chain with the `else`
+          // becoming part of the next section's header.
+          traverse(ifElement.elseKeyword, elseIf);
+
+        case var elseElement?:
+          var spreadElse = elseElement.spreadCollection;
+
+          // Any other kind of else body ends the chain, with the header for
+          // the last section just being the `else` keyword.
+          var header = buildPiece((b) {
+            b.token(ifElement.elseKeyword!);
+
+            // Make the `...` part of the header so that IfPiece can correctly
+            // constrain the inner collection literal's ListPiece to split.
+            if (spreadElse != null) {
+              b.space();
+              b.token(spreadElse.spreadOperator);
+            }
+          });
+
+          Piece statement;
+          if (spreadElse != null) {
+            statement = nodePiece(spreadElse.expression);
+          } else {
+            statement = nodePiece(elseElement);
+          }
+
+          piece.add(header, statement, isBlock: spreadElse != null);
+
+          // If the else branch of an if element is itself a control flow
+          // element, then force the outer if to always split.
+          if (ifElement.thenElement.isControlFlowElement) {
+            piece.pin(State.split);
+          }
+
+        case null:
+          break; // Nothing to do.
+      }
+    }
+
+    traverse(null, node);
+    return piece;
   }
 
   @override
   Piece visitIfStatement(IfStatement node) {
-    return createIf(node);
+    var piece = IfPiece();
+
+    // Recurses through the else branches to flatten them into a linear if-else
+    // chain handled by a single [IfPiece].
+    void traverse(Token? precedingElse, IfStatement ifStatement) {
+      var condition = buildPiece((b) {
+        b.token(precedingElse, spaceAfter: true);
+        b.add(startControlFlow(
+            ifStatement.ifKeyword,
+            ifStatement.leftParenthesis,
+            ifStatement.expression,
+            ifStatement.rightParenthesis));
+        b.space();
+      });
+
+      // Edge case: When the then branch is a block and there is an else clause
+      // after it, we want to force the block to split even if empty, like:
+      //
+      // ```
+      // if (condition) {
+      // } else {
+      //   body;
+      // }
+      // ```
+      var thenStatement = switch (ifStatement.thenStatement) {
+        Block thenBlock when ifStatement.elseStatement != null =>
+          createBlock(thenBlock, forceSplit: true),
+        _ => nodePiece(ifStatement.thenStatement)
+      };
+
+      piece.add(condition, thenStatement,
+          isBlock: ifStatement.thenStatement is Block);
+
+      switch (ifStatement.elseStatement) {
+        case IfStatement elseIf:
+          // Hit an else-if, so flatten it into the chain with the `else`
+          // becoming part of the next section's header.
+          traverse(ifStatement.elseKeyword, elseIf);
+
+        case var elseStatement?:
+          // Any other kind of else body ends the chain, with the header for
+          // the last section just being the `else` keyword.
+          var header = buildPiece((b) {
+            b.token(ifStatement.elseKeyword, spaceAfter: true);
+          });
+          var statement = nodePiece(elseStatement);
+          piece.add(header, statement, isBlock: elseStatement is Block);
+      }
+    }
+
+    traverse(null, node);
+
+    // If statements almost always split at the clauses unless the if is a
+    // simple if with only a single unbraced then statement and no else clause,
+    // like:
+    //
+    //     if (condition) print("ok");
+    if (node.thenStatement is Block || node.elseStatement != null) {
+      piece.pin(State.split);
+    }
+
+    return piece;
   }
 
   @override
@@ -1666,8 +1805,11 @@ class AstNodeVisitor extends ThrowingAstVisitor<Piece> with PieceFactory {
 
   @override
   Piece visitWhileStatement(WhileStatement node) {
-    var condition = startControlFlow(node.whileKeyword, node.leftParenthesis,
-        node.condition, node.rightParenthesis);
+    var condition = buildPiece((b) {
+      b.add(startControlFlow(node.whileKeyword, node.leftParenthesis,
+          node.condition, node.rightParenthesis));
+      b.space();
+    });
 
     var body = nodePiece(node.body);
 

lib/src/front_end/piece_factory.dart

@@ -9,7 +9,6 @@ import '../piece/assign.dart';
 import '../piece/block.dart';
 import '../piece/clause.dart';
 import '../piece/function.dart';
-import '../piece/if.dart';
 import '../piece/infix.dart';
 import '../piece/list.dart';
 import '../piece/piece.dart';
@@ -405,59 +404,6 @@ mixin PieceFactory {
     return piece;
   }
 
-  // TODO(tall): Generalize this to work with if elements too.
-  /// Creates a piece for a chain of if-else-if... statements.
-  Piece createIf(IfStatement ifStatement) {
-    var piece = IfPiece();
-
-    // Recurses through the else branches to flatten them into a linear if-else
-    // chain handled by a single [IfPiece].
-    void traverse(Piece? previousElse, IfStatement node) {
-      var condition = buildPiece((b) {
-        if (previousElse != null) b.add(previousElse);
-        b.add(startControlFlow(node.ifKeyword, node.leftParenthesis,
-            node.expression, node.rightParenthesis));
-      });
-
-      // Edge case: When the then branch is a block and there is an else clause
-      // after it, we want to force the block to split even if empty, like:
-      //
-      // ```
-      // if (condition) {
-      // } else {
-      //   body;
-      // }
-      // ```
-      var thenStatement = switch (node.thenStatement) {
-        Block thenBlock when node.elseStatement != null =>
-          createBlock(thenBlock, forceSplit: true),
-        _ => nodePiece(node.thenStatement)
-      };
-
-      piece.add(condition, thenStatement, isBlock: node.thenStatement is Block);
-
-      switch (node.elseStatement) {
-        case IfStatement elseIf:
-          // Hit an else-if, so flatten it into the chain with the `else`
-          // becoming part of the next section's header.
-          traverse(buildPiece((b) {
-            b.token(node.elseKeyword);
-            b.space();
-          }), elseIf);
-
-        case var elseStatement?:
-          // Any other kind of else body ends the chain, with the header for
-          // the last section just being the `else` keyword.
-          var header = tokenPiece(node.elseKeyword!);
-          var statement = nodePiece(elseStatement);
-          piece.add(header, statement, isBlock: elseStatement is Block);
-      }
-    }
-
-    traverse(null, ifStatement);
-    return piece;
-  }
-
   /// Creates an [ImportPiece] for an import or export directive.
   Piece createImport(NamespaceDirective directive, Token keyword,
       {Token? deferredKeyword, Token? asKeyword, SimpleIdentifier? prefix}) {

lib/src/piece/if.dart

@@ -5,76 +5,64 @@ import '../back_end/code_writer.dart';
 import '../constants.dart';
 import 'piece.dart';
 
-/// A piece for an if statement.
+/// A piece for an if statement or element.
 ///
 /// We also use this for while statements, which are formatted exactly like an
 /// if statement with no else clause.
 class IfPiece extends Piece {
   final List<_IfSection> _sections = [];
 
-  /// Whether the if is a simple if with only a single unbraced then statement
-  /// and no else clause, like:
-  ///
-  /// ```
-  /// if (condition) print("ok");
-  /// ```
-  ///
-  /// Unlike other if statements, these allow a discretionary split after the
-  /// condition.
-  bool get _isUnbracedIfThen =>
-      _sections.length == 1 && !_sections.single.isBlock;
-
   void add(Piece header, Piece statement, {required bool isBlock}) {
     _sections.add(_IfSection(header, statement, isBlock));
   }
 
-  /// If there is at least one else or else-if clause, then it always splits.
   @override
-  List<State> get additionalStates => [if (_isUnbracedIfThen) State.split];
+  List<State> get additionalStates => [State.split];
 
   @override
-  void format(CodeWriter writer, State state) {
-    if (_isUnbracedIfThen) {
-      // A split in the condition or statement forces moving the entire
-      // statement to the next line.
-      writer.setAllowNewlines(state != State.unsplit);
+  void applyConstraints(State state, Constrain constrain) {
+    // 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);
+      }
+    }
+  }
 
-      var section = _sections.single;
-      writer.format(section.header);
-      writer.splitIf(state == State.split, indent: Indent.block);
-      writer.format(section.statement);
-      return;
+  @override
+  void forEachChild(void Function(Piece piece) callback) {
+    for (var section in _sections) {
+      callback(section.header);
+      callback(section.statement);
     }
+  }
 
+  @override
+  void format(CodeWriter writer, State state) {
     for (var i = 0; i < _sections.length; i++) {
       var section = _sections[i];
 
+      // A split in the condition forces the branches to split.
+      writer.setAllowNewlines(state == State.split);
       writer.format(section.header);
 
-      // If the statement is a block, then keep the `{` on the same line as the
-      // header part.
-      if (section.isBlock) {
-        writer.space();
-      } else {
-        writer.newline(indent: Indent.block);
+      if (!section.isBlock) {
+        writer.splitIf(state == State.split, indent: Indent.block);
       }
 
       writer.format(section.statement);
 
       // Reset the indentation for the subsequent `else` or `} else` line.
       if (i < _sections.length - 1) {
-        writer.splitIf(!section.isBlock, indent: Indent.none);
+        writer.splitIf(state == State.split && !section.isBlock,
+            indent: Indent.none);
       }
     }
   }
-
-  @override
-  void forEachChild(void Function(Piece piece) callback) {
-    for (var section in _sections) {
-      callback(section.header);
-      callback(section.statement);
-    }
-  }
 }
 
 /// A single section in a chain of if-elses.
@@ -88,7 +76,8 @@ class _IfSection {
   final Piece header;
   final Piece statement;
 
-  /// Whether the [statement] piece is from a block.
+  /// Whether the [statement] piece is from a block or a spread collection
+  /// literal.
   final bool isBlock;
 
   _IfSection(this.header, this.statement, this.isBlock);