Format if and return statements. (#1284)

* Format if and return statements.

I also tweaked the CodeWriter API to make it easier to update the
indentation when inserting a newline. This doesn't do anything radically
new, but makes the piece formatting code a little shorter.

Made a field in SequencePiece private that didn't need to be public.

* Fix class comment.

Author: munificent

lib/src/back_end/code_writer.dart

@@ -141,8 +141,13 @@ class CodeWriter {
 
   /// Inserts a line split in the output.
   ///
-  /// If [blank] is true, writes an extra newline to produce a blank line.
-  void newline({bool blank = false}) {
+  /// If [blank] is `true`, writes an extra newline to produce a blank line.
+  ///
+  /// If [indent] is given, set the indentation of the new line (and all
+  /// subsequent lines) to that indentation relative to the containing piece.
+  void newline({bool blank = false, int? indent}) {
+    if (indent != null) setIndent(indent);
+
     handleNewline();
     _finishLine();
     _buffer.writeln();

lib/src/front_end/ast_node_visitor.dart

@@ -143,7 +143,7 @@ class AstNodeVisitor extends ThrowingAstVisitor<void>
 
   @override
   void visitBlock(Block node) {
-    createBlock(node.leftBracket, node.statements, node.rightBracket);
+    createBlock(node);
   }
 
   @override
@@ -441,7 +441,7 @@ class AstNodeVisitor extends ThrowingAstVisitor<void>
 
   @override
   void visitIfStatement(IfStatement node) {
-    throw UnimplementedError();
+    createIf(node);
   }
 
   @override
@@ -766,7 +766,14 @@ class AstNodeVisitor extends ThrowingAstVisitor<void>
 
   @override
   void visitReturnStatement(ReturnStatement node) {
-    throw UnimplementedError();
+    token(node.returnKeyword);
+
+    if (node.expression case var expression) {
+      writer.space();
+      visit(expression);
+    }
+
+    token(node.semicolon);
   }
 
   @override

lib/src/front_end/piece_factory.dart

@@ -7,6 +7,7 @@ import 'package:analyzer/dart/ast/token.dart';
 import '../ast_extensions.dart';
 import '../piece/assign.dart';
 import '../piece/block.dart';
+import '../piece/if.dart';
 import '../piece/import.dart';
 import '../piece/infix.dart';
 import '../piece/piece.dart';
@@ -46,33 +47,42 @@ mixin PieceFactory implements CommentWriter {
 
   /// Creates a [BlockPiece] for a given bracket-delimited block or declaration
   /// body.
-  void createBlock(Token leftBracket, List<AstNode> nodes, Token rightBracket) {
+  ///
+  /// If [forceSplit] is `true`, then the block will split even if empty. This
+  /// is used, for example, with empty blocks in `if` statements followed by
+  /// `else` clauses:
+  ///
+  /// ```
+  /// if (condition) {
+  /// } else {}
+  /// ```
+  void createBlock(Block block, {bool forceSplit = false}) {
     // Edge case: If the block is completely empty, output it as simple
-    // unsplittable text.
-    if (nodes.isEmptyBody(rightBracket)) {
-      token(leftBracket);
-      token(rightBracket);
+    // unsplittable text unless it's forced to split.
+    if (block.statements.isEmptyBody(block.rightBracket) && !forceSplit) {
+      token(block.leftBracket);
+      token(block.rightBracket);
       return;
     }
 
-    token(leftBracket);
+    token(block.leftBracket);
     var leftBracketPiece = writer.pop();
     writer.split();
 
     var sequence = SequenceBuilder(this);
-    for (var node in nodes) {
+    for (var node in block.statements) {
       sequence.add(node);
     }
 
     // Place any comments before the "}" inside the block.
-    sequence.addCommentsBefore(rightBracket);
+    sequence.addCommentsBefore(block.rightBracket);
 
-    token(rightBracket);
+    token(block.rightBracket);
     var rightBracketPiece = writer.pop();
 
     writer.push(BlockPiece(
         leftBracketPiece, sequence.build(), rightBracketPiece,
-        alwaysSplit: nodes.isNotEmpty));
+        alwaysSplit: forceSplit || block.statements.isNotEmpty));
   }
 
   /// Creates a [ListPiece] for a collection literal.
@@ -123,6 +133,69 @@ mixin PieceFactory implements CommentWriter {
     }
   }
 
+  // TODO(tall): Generalize this to work with if elements too.
+  /// Creates a piece for a chain of if-else-if... statements.
+  void 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(IfStatement node) {
+      token(node.ifKeyword);
+      writer.space();
+      token(node.leftParenthesis);
+      visit(node.expression);
+      token(node.rightParenthesis);
+      var condition = writer.pop();
+      writer.split();
+
+      // 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;
+      // }
+      // ```
+      if (node.thenStatement case Block thenBlock
+          when node.elseStatement != null) {
+        createBlock(thenBlock, forceSplit: true);
+      } else {
+        visit(node.thenStatement);
+      }
+
+      var thenStatement = writer.pop();
+      writer.split();
+      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.
+          token(node.elseKeyword);
+          writer.space();
+          traverse(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.
+          token(node.elseKeyword);
+          var header = writer.pop();
+          writer.split();
+
+          visit(elseStatement);
+          var statement = writer.pop();
+          writer.split();
+          piece.add(header, statement, isBlock: elseStatement is Block);
+      }
+    }
+
+    traverse(ifStatement);
+
+    writer.push(piece);
+  }
+
   /// Creates an [ImportPiece] for an import or export directive.
   void createImport(NamespaceDirective directive, Token keyword,
       {Token? deferredKeyword, Token? asKeyword, SimpleIdentifier? prefix}) {

lib/src/piece/assign.dart

@@ -72,7 +72,6 @@ class AssignPiece extends Piece {
     if (state != _insideValue) writer.setIndent(Indent.expression);
 
     writer.format(target);
-
     writer.splitIf(state == _atEquals);
     writer.format(value);
   }

lib/src/piece/block.dart

@@ -37,11 +37,12 @@ class BlockPiece extends Piece {
     writer.format(leftBracket);
 
     if (_alwaysSplit || state == State.split) {
-      writer.setIndent(Indent.block);
-      writer.newline();
-      writer.format(contents);
-      writer.setIndent(Indent.none);
-      writer.newline();
+      if (contents.isNotEmpty) {
+        writer.newline(indent: Indent.block);
+        writer.format(contents);
+      }
+
+      writer.newline(indent: Indent.none);
     } else {
       writer.setAllowNewlines(false);
       writer.format(contents);

lib/src/piece/if.dart

@@ -0,0 +1,95 @@
+// Copyright (c) 2023, 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 '../constants.dart';
+import 'piece.dart';
+
+/// A piece for an if statement.
+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 states => _isUnbracedIfThen ? const [State.split] : const [];
+
+  @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.initial);
+
+      var section = _sections.single;
+      writer.format(section.header);
+      writer.splitIf(state == State.split, indent: Indent.block);
+      writer.format(section.statement);
+      return;
+    }
+
+    for (var i = 0; i < _sections.length; i++) {
+      var section = _sections[i];
+
+      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);
+      }
+
+      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);
+      }
+    }
+  }
+
+  @override
+  void forEachChild(void Function(Piece piece) callback) {
+    for (var section in _sections) {
+      callback(section.header);
+      callback(section.statement);
+    }
+  }
+
+  @override
+  String toString() => 'If';
+}
+
+/// A single section in a chain of if-elses.
+///
+/// For the first then branch, the [header] is the `if (condition)` part and
+/// the statement is the then branch. For all `else if` branches, the [header]
+/// is the `else if (condition)` and the statement is the subsequent then
+/// branch. For the final `else` branch, if there is one, the [header] is just
+/// `else` and the statement is the else branch.
+class _IfSection {
+  final Piece header;
+  final Piece statement;
+
+  /// Whether the [statement] piece is from a block.
+  final bool isBlock;
+
+  _IfSection(this.header, this.statement, this.isBlock);
+}

lib/src/piece/list.dart

@@ -70,8 +70,7 @@ class ListPiece extends Piece {
 
       case State.split:
         // Each argument on its own line with a trailing comma after the last.
-        writer.setIndent(Indent.block);
-        writer.newline();
+        writer.newline(indent: Indent.block);
         for (var i = 0; i < _arguments.length; i++) {
           var argument = _arguments[i];
           argument.format(writer,
@@ -80,8 +79,7 @@ class ListPiece extends Piece {
             writer.newline(blank: _blanksAfter.contains(argument));
           }
         }
-        writer.setIndent(Indent.none);
-        writer.newline();
+        writer.newline(indent: Indent.none);
     }
 
     writer.setAllowNewlines(true);

lib/src/piece/sequence.dart

@@ -11,31 +11,34 @@ import 'piece.dart';
 /// Usually constructed using a [SequenceBuilder].
 class SequencePiece extends Piece {
   /// The series of members or statements.
-  final List<Piece> contents;
+  final List<Piece> _contents;
 
   /// The pieces that should have a blank line preserved between them and the
   /// next piece.
   final Set<Piece> _blanksAfter;
 
-  SequencePiece(this.contents, this._blanksAfter);
+  SequencePiece(this._contents, this._blanksAfter);
+
+  /// Whether this sequence has any contents.
+  bool get isNotEmpty => _contents.isNotEmpty;
 
   @override
   List<State> get states => const [];
 
   @override
   void format(CodeWriter writer, State state) {
-    for (var i = 0; i < contents.length; i++) {
-      writer.format(contents[i]);
+    for (var i = 0; i < _contents.length; i++) {
+      writer.format(_contents[i]);
 
-      if (i < contents.length - 1) {
-        writer.newline(blank: _blanksAfter.contains(contents[i]));
+      if (i < _contents.length - 1) {
+        writer.newline(blank: _blanksAfter.contains(_contents[i]));
       }
     }
   }
 
   @override
   void forEachChild(void Function(Piece piece) callback) {
-    contents.forEach(callback);
+    _contents.forEach(callback);
   }
 
   @override