From 729224abf615f866a3a4e6e99fb6792f0ea4950b Mon Sep 17 00:00:00 2001
From: Robert Nystrom <rnystrom@google.com>
Date: Tue, 14 Nov 2023 17:16:28 -0800
Subject: [PATCH 1/3] Support "block argument" formatting.

Sometimes in an argument list, we don't want to force it to full split
even if an argument contains an inner newline. In other words, we want
to prefer:

```
test('hi', () {
  body;
});
```

Over:

```
test(
  'hi',
  () {
    body;
  },
);
```

The latter is what you would normally get when an argument contains a
newline. This change enables the former. The heuristics for what kinds
of arguments are allowed to have this block-like formatting, how many
of them, and what other arguments are allowed are pretty fuzzy.

I picked a simple set of rules here (which follow the prototype), but
left some TODOs because I suspect we will iterate on them later once we
can run the new formatter on larger codebases.

The heuristic for what kinds of arguments are block-like is the same
rule that we use for "delimited" expressions after an assignment. I
think it's good that those two agree and helps make the style overall
more consistent. Basically, any time the right-hand side of an `=` or
`:` doesn't need to split, that same kind of expression can be block
formatted:

```
var list = [
  element,
];

function([
  element,
]);
```

Block formatting is used for argument lists, asserts, and switch values.

I also simplified AssignPiece. It used an extra unnecessary to allow
newlines inside delimited value expressions without splitting after "=".
---
 lib/src/ast_extensions.dart                   |  66 +++-
 lib/src/front_end/ast_node_visitor.dart       |  28 +-
 lib/src/front_end/comment_writer.dart         |   6 +-
 lib/src/front_end/delimited_list_builder.dart |  62 +++-
 lib/src/front_end/piece_factory.dart          |  26 +-
 lib/src/piece/assign.dart                     |  27 +-
 lib/src/piece/list.dart                       |  78 ++++-
 test/expression/switch.stmt                   |  18 +-
 test/invocation/block_argument_comment.stmt   | 107 +++++++
 test/invocation/block_argument_kind.stmt      | 284 ++++++++++++++++++
 test/invocation/block_argument_multiple.stmt  |  66 ++++
 test/invocation/block_argument_other.stmt     |  89 ++++++
 test/invocation/constructor.stmt              |  10 +-
 test/invocation/function.stmt                 |  23 +-
 test/statement/assert.stmt                    |  18 +-
 test/statement/switch.stmt                    |  12 +-
 test/variable/local.stmt                      |  18 +-
 17 files changed, 830 insertions(+), 108 deletions(-)
 create mode 100644 test/invocation/block_argument_comment.stmt
 create mode 100644 test/invocation/block_argument_kind.stmt
 create mode 100644 test/invocation/block_argument_multiple.stmt
 create mode 100644 test/invocation/block_argument_other.stmt

diff --git a/lib/src/ast_extensions.dart b/lib/src/ast_extensions.dart
index a6cc3cde..265e00cc 100644
--- a/lib/src/ast_extensions.dart
+++ b/lib/src/ast_extensions.dart
@@ -126,17 +126,61 @@ extension ExpressionExtensions on Expression {
   ///   element,
   /// ];
   /// ```
-  bool get isDelimited => switch (this) {
-        FunctionExpression() => true,
-        InstanceCreationExpression() => true,
-        ListLiteral() => true,
-        MethodInvocation() => true,
-        ParenthesizedExpression(:var expression) => expression.isDelimited,
-        RecordLiteral() => true,
-        SetOrMapLiteral() => true,
-        SwitchExpression() => true,
-        _ => false,
-      };
+  ///
+  /// Likewise, in an argument list, delimited expressions can be given
+  /// block-like formatting:
+  ///
+  /// ```
+  /// function([
+  ///   element,
+  /// ]);
+  /// ```
+  bool get isDelimited {
+    // Unwrap named expressions to get the real expression inside.
+    var expression = this;
+    if (expression is NamedExpression) {
+      expression = expression.expression;
+    }
+
+    // TODO(tall): We should also allow multi-line strings to be formatted
+    // like block arguments, at least in some cases like:
+    //
+    // ```
+    // function('''
+    //   Lots of
+    //   text
+    // ''');
+    // ```
+
+    // TODO(tall): Consider whether immediately-invoked function expressions
+    // should be block argument candidates, like:
+    //
+    // ```
+    // function(() {
+    //   body;
+    // }());
+    // ```
+    return switch (expression) {
+      // A function expression can use either a non-empty parameter list or a
+      // non-empty block body for block formatting.
+      FunctionExpression(:var parameters!, :var body) =>
+        !parameters.parameters.isEmptyBody(parameters.rightParenthesis) ||
+            body is BlockFunctionBody &&
+                !body.block.statements.isEmptyBody(body.block.rightBracket),
+      ListLiteral(:var elements, :var rightBracket) ||
+      SetOrMapLiteral(:var elements, :var rightBracket) =>
+        !elements.isEmptyBody(rightBracket),
+      RecordLiteral(:var fields, :var rightParenthesis) =>
+        !fields.isEmptyBody(rightParenthesis),
+      SwitchExpression(:var cases, :var rightBracket) =>
+        !cases.isEmptyBody(rightBracket),
+      InstanceCreationExpression(:var argumentList) ||
+      MethodInvocation(:var argumentList) =>
+        !argumentList.arguments.isEmptyBody(argumentList.rightParenthesis),
+      ParenthesizedExpression(:var expression) => expression.isDelimited,
+      _ => false,
+    };
+  }
 
   /// Whether this is an argument in an argument list with a trailing comma.
   bool get isTrailingCommaArgument {
diff --git a/lib/src/front_end/ast_node_visitor.dart b/lib/src/front_end/ast_node_visitor.dart
index 8ea802aa..7700e0a1 100644
--- a/lib/src/front_end/ast_node_visitor.dart
+++ b/lib/src/front_end/ast_node_visitor.dart
@@ -115,10 +115,8 @@ class AstNodeVisitor extends ThrowingAstVisitor<void>
 
   @override
   void visitArgumentList(ArgumentList node) {
-    createList(
-        leftBracket: node.leftParenthesis,
-        node.arguments,
-        rightBracket: node.rightParenthesis);
+    createArgumentList(
+        node.leftParenthesis, node.arguments, node.rightParenthesis);
   }
 
   @override
@@ -134,14 +132,13 @@ class AstNodeVisitor extends ThrowingAstVisitor<void>
   @override
   void visitAssertStatement(AssertStatement node) {
     token(node.assertKeyword);
-    createList(
-      [
-        node.condition,
-        if (node.message case var message?) message,
-      ],
-      leftBracket: node.leftParenthesis,
-      rightBracket: node.rightParenthesis,
-    );
+    createArgumentList(
+        node.leftParenthesis,
+        [
+          node.condition,
+          if (node.message case var message?) message,
+        ],
+        node.rightParenthesis);
     token(node.semicolon);
   }
 
@@ -438,7 +435,6 @@ class AstNodeVisitor extends ThrowingAstVisitor<void>
               spaceWhenUnsplit: true, splitListIfBeforeSplits: true));
       builder.leftBracket(node.leftBracket);
       node.constants.forEach(builder.visit);
-
       builder.rightBracket(semicolon: node.semicolon, node.rightBracket);
       pieces.give(builder.build());
     } else {
@@ -959,8 +955,10 @@ class AstNodeVisitor extends ThrowingAstVisitor<void>
 
   @override
   void visitMethodInvocation(MethodInvocation node) {
-    // TODO(tall): Support method invocation with explicit target expressions.
-    if (node.target != null) throw UnimplementedError();
+    // TODO(tall): Support splitting at `.` or `?.`. Right now we just format
+    // it inline so that we can use method calls in other tests.
+    visit(node.target);
+    token(node.operator);
 
     visit(node.methodName);
     visit(node.typeArguments);
diff --git a/lib/src/front_end/comment_writer.dart b/lib/src/front_end/comment_writer.dart
index a7d5a733..1e9c68a8 100644
--- a/lib/src/front_end/comment_writer.dart
+++ b/lib/src/front_end/comment_writer.dart
@@ -34,14 +34,14 @@ import 'piece_writer.dart';
 /// couple of different ways.
 ///
 /// Comments between top-level declarations, member declarations inside types,
-/// and statements are handled directly by [SequenceBuilder].
+/// and statements are handled directly by [SequenceBuilder]. Comments inside
+/// argument lists, collection literals, and other similar constructs are
+/// handled directly be [DelimitedPieceBuilder].
 ///
 /// All other comments occur inside the middle of some expression or other
 /// construct. These get directly embedded in the [TextPiece] of the code being
 /// written. When that [TextPiece] is output later, it will include the comments
 /// as well.
-// TODO(tall): When argument lists and their comment handling is supported,
-// mention that here.
 mixin CommentWriter {
   PieceWriter get pieces;
 
diff --git a/lib/src/front_end/delimited_list_builder.dart b/lib/src/front_end/delimited_list_builder.dart
index d2e8cfbb..b7940590 100644
--- a/lib/src/front_end/delimited_list_builder.dart
+++ b/lib/src/front_end/delimited_list_builder.dart
@@ -36,6 +36,11 @@ class DelimitedListBuilder {
 
   final ListStyle _style;
 
+  /// The nodes that have been visited by this builder.
+  ///
+  /// Used to find the block element from them, if desired.
+  final List<AstNode> _blockCandidates = [];
+
   /// The list of comments following the most recently written element before
   /// any comma following the element.
   CommentSequence _commentsBeforeComma = CommentSequence.empty;
@@ -44,8 +49,15 @@ class DelimitedListBuilder {
   /// literal, etc.
   DelimitedListBuilder(this._visitor, [this._style = const ListStyle()]);
 
-  ListPiece build() =>
-      ListPiece(_leftBracket, _elements, _blanksAfter, _rightBracket, _style);
+  /// Create the final [ListPiece] out of the added brackets, delimiters,
+  /// elements, and style.
+  ListPiece build() {
+    var blockElement = -1;
+    if (_style.allowBlockElement) blockElement = _findBlockElement();
+
+    return ListPiece(_leftBracket, _elements, _blanksAfter, _rightBracket,
+        _style, blockElement);
+  }
 
   /// Adds the opening [bracket] to the built list.
   ///
@@ -144,6 +156,9 @@ class DelimitedListBuilder {
     if (nextToken.lexeme == ',') {
       _commentsBeforeComma = _visitor.takeCommentsBefore(nextToken);
     }
+
+    // If this element might be block formatted, remember it.
+    if (_style.allowBlockElement) _blockCandidates.add(element);
   }
 
   /// Inserts an inner left delimiter between two elements.
@@ -373,4 +388,47 @@ class DelimitedListBuilder {
       leading: leadingComments
     );
   }
+
+  /// If [_blockCandidates] contains a single expression that can receive
+  /// block formatting, then returns its index. Otherwise returns `-1`.
+  int _findBlockElement() {
+    // TODO(tall): These heuristics will probably need some iteration.
+    var functions = <int>[];
+    var others = <int>[];
+
+    for (var i = 0; i < _blockCandidates.length; i++) {
+      var node = _blockCandidates[i];
+      if (node is Expression && node.isDelimited) {
+        if (node is FunctionExpression) {
+          functions.add(i);
+        } else {
+          others.add(i);
+        }
+      }
+    }
+
+    // A function expression takes precedence over other block arguments.
+    if (functions.length == 1) return functions.first;
+
+    // Otherwise, if there is single block argument, it can be block formatted.
+    if (functions.isEmpty && others.length == 1) return others.first;
+
+    // There are no block arguments, or it's ambiguous as to which one should
+    // be it.
+    // TODO(tall): The old formatter allows multiple block arguments, like:
+    //
+    // ```
+    // function(() {
+    //   body;
+    // }, () {
+    //   more;
+    // });
+    // ```
+    //
+    // This doesn't seem very common in the Flutter repo, but does occur
+    // sometimes. We'll probably want to experiment to see if it's worth
+    // supporting multiple block arguments. If so, we should at least require
+    // them to be contiguous with no non-block arguments in the middle.
+    return -1;
+  }
 }
diff --git a/lib/src/front_end/piece_factory.dart b/lib/src/front_end/piece_factory.dart
index ece6f697..a4e8f323 100644
--- a/lib/src/front_end/piece_factory.dart
+++ b/lib/src/front_end/piece_factory.dart
@@ -49,6 +49,16 @@ typedef BinaryOperation = (AstNode left, Token operator, AstNode right);
 mixin PieceFactory implements CommentWriter {
   void visit(AstNode? node, {void Function()? before, void Function()? after});
 
+  /// Creates a [ListPiece] for an argument list.
+  void createArgumentList(
+      Token leftBracket, Iterable<AstNode> elements, Token rightBracket) {
+    return createList(
+        leftBracket: leftBracket,
+        elements,
+        rightBracket: rightBracket,
+        style: const ListStyle(allowBlockElement: true));
+  }
+
   /// Creates a [BlockPiece] for a given bracket-delimited block or declaration
   /// body.
   ///
@@ -425,20 +435,16 @@ mixin PieceFactory implements CommentWriter {
   /// Visits the `switch (expr)` part of a switch statement or expression.
   void createSwitchValue(Token switchKeyword, Token leftParenthesis,
       Expression value, Token rightParenthesis) {
-    // Format like an argument list since it is an expression surrounded by
-    // parentheses.
-    var builder = DelimitedListBuilder(
-        this, const ListStyle(commas: Commas.none, splitCost: 2));
-
     // Attach the `switch ` as part of the `(`.
     token(switchKeyword);
     space();
 
-    builder.leftBracket(leftParenthesis);
-    builder.visit(value);
-    builder.rightBracket(rightParenthesis);
-
-    pieces.give(builder.build());
+    createList(
+        leftBracket: leftParenthesis,
+        [value],
+        rightBracket: rightParenthesis,
+        style: const ListStyle(
+            commas: Commas.none, splitCost: 2, allowBlockElement: true));
   }
 
   /// Creates a class, enum, extension, mixin, or mixin application class
diff --git a/lib/src/piece/assign.dart b/lib/src/piece/assign.dart
index a17ac862..c949dabc 100644
--- a/lib/src/piece/assign.dart
+++ b/lib/src/piece/assign.dart
@@ -21,9 +21,8 @@ import 'piece.dart';
 /// var x = 123;
 /// ```
 ///
-/// [_insideValue] If the value is a delimited "block-like" expression,
-/// then we can split inside the block but not at the `=` with no additional
-/// indentation:
+/// If the value is a delimited "block-like" expression, then we allow splitting
+/// inside the value but not at the `=` with no additional indentation:
 ///
 /// ```
 /// var list = [
@@ -31,22 +30,17 @@ import 'piece.dart';
 /// ];
 /// ```
 ///
-/// [State.split] Split after the `=`:
+/// [_atOperator] Split after the `=`:
 ///
 /// ```
 /// var name =
 ///     longValueExpression;
 /// ```
 class AssignPiece extends Piece {
-  /// Split inside the value but not at the `=`.
+  /// Split after the operator.
   ///
-  /// This is only allowed when the value is a delimited expression.
-  static const State _insideValue = State(1);
-
-  /// Split at the operator and allow splitting inside the value.
-  ///
-  /// This is more costly because, when it's possible to split inside a
-  /// delimited value, we want to prefer that.
+  /// This is more costly because it's generally better to split either in the
+  /// value (if it's delimited) or in the target.
   static const State _atOperator = State(2, cost: 2);
 
   /// The left-hand side of the operation. Includes the operator unless it is
@@ -107,17 +101,18 @@ class AssignPiece extends Piece {
   // ```
 
   @override
-  List<State> get additionalStates =>
-      [if (_isValueDelimited) _insideValue, _atOperator];
+  List<State> get additionalStates => [_atOperator];
 
   @override
   void format(CodeWriter writer, State state) {
     // A split in either child piece forces splitting after the "=" unless it's
     // a delimited expression.
-    if (state == State.unsplit) writer.setAllowNewlines(false);
+    if (state == State.unsplit && !_isValueDelimited) {
+      writer.setAllowNewlines(false);
+    }
 
     // Don't indent a split delimited expression.
-    if (state != _insideValue) writer.setIndent(Indent.expression);
+    if (state != State.unsplit) writer.setIndent(Indent.expression);
 
     writer.format(target);
     writer.splitIf(state == _atOperator);
diff --git a/lib/src/piece/list.dart b/lib/src/piece/list.dart
index b5b9d2ff..5713eb15 100644
--- a/lib/src/piece/list.dart
+++ b/lib/src/piece/list.dart
@@ -15,7 +15,35 @@ import 'piece.dart';
 /// the list is split or not. It handles comments inside the sequence of
 /// elements.
 ///
-/// Usually constructed using [createList()] or a [DelimitedListBuilder].
+/// These pieces can be formatted in one of three ways:
+///
+/// [State.split] Fully unsplit:
+///
+/// ```
+/// function(argument, argument, argument);
+/// ```
+///
+/// If one of the elements is a "block element", then we allow newlines inside
+/// it to support output like:
+///
+/// ```
+/// function(argument, () {
+///   blockElement;
+/// }, argument);
+/// ```
+///
+/// [_splitState] Split around all of the items:
+///
+/// ```
+/// function(
+///   argument,
+///   argument,
+///   argument,
+/// );
+/// ```
+///
+/// ListPieces are usually constructed using [createList()] or
+/// [DelimitedListBuilder].
 class ListPiece extends Piece {
   /// The opening bracket before the elements, if any.
   final Piece? _before;
@@ -37,29 +65,28 @@ class ListPiece extends Piece {
   ///
   /// We use this instead of [State.split] because the cost is higher for some
   /// kinds of lists.
+  ///
+  /// Also, if this list does have a block element, then we also increase the
+  /// cost of splitting the list itself a little more to prefer the block
+  /// argument splitting when possible.
   // TODO(rnystrom): Having to use a different state for this is a little
   // cumbersome. Maybe it would be better to most costs out of [State] and
   // instead have the [Solver] ask each [Piece] for the cost of its state.
   final State _splitState;
 
-  ListPiece(
-      this._before, this._elements, this._blanksAfter, this._after, this._style)
-      : _splitState = State(1, cost: _style.splitCost);
+  /// If this list has an element that can receive block formatting, this is
+  /// the elements's index. Otherwise `-1`.
+  final int _blockElement;
+
+  ListPiece(this._before, this._elements, this._blanksAfter, this._after,
+      this._style, this._blockElement)
+      : _splitState = State(2, cost: _style.splitCost);
 
   @override
   List<State> get additionalStates => [if (_elements.isNotEmpty) _splitState];
 
   @override
   void format(CodeWriter writer, State state) {
-    // TODO(tall): Should support a third state for argument lists with block
-    // arguments, like:
-    //
-    // ```
-    // test('description', () {
-    //   ...
-    // });
-    // ```
-
     // Format the opening bracket, if there is one.
     if (_before case var before?) {
       if (_style.splitListIfBeforeSplits && state == State.unsplit) {
@@ -89,9 +116,18 @@ class ListPiece extends Piece {
         Commas.none => false,
       };
 
+      // Only allow newlines in the block element or in all elements if we're
+      // fully split.
+      writer.setAllowNewlines(i == _blockElement || state == _splitState);
+
       var element = _elements[i];
       element.format(writer, appendComma: appendComma);
 
+      // Only allow newlines in comments if we're fully split.
+      writer.setAllowNewlines(state == _splitState);
+
+      element.formatComment(writer);
+
       // Write a space or newline between elements.
       if (!isLast) {
         writer.splitIf(state != State.unsplit,
@@ -186,7 +222,9 @@ final class ListElement {
         writer.write(_delimiter);
       }
     }
+  }
 
+  void formatComment(CodeWriter writer) {
     if (_comment case var comment?) {
       if (_element != null) writer.space();
       writer.format(comment);
@@ -293,9 +331,21 @@ class ListStyle {
   /// ```
   final bool splitListIfBeforeSplits;
 
+  /// Whether an element in the list is allowed to have block-like formatting,
+  /// as in:
+  ///
+  /// ```
+  /// function(argument, [
+  ///   block,
+  ///   like,
+  /// ], argument);
+  /// ```
+  final bool allowBlockElement;
+
   const ListStyle(
       {this.commas = Commas.trailing,
       this.splitCost = Cost.normal,
       this.spaceWhenUnsplit = false,
-      this.splitListIfBeforeSplits = false});
+      this.splitListIfBeforeSplits = false,
+      this.allowBlockElement = false});
 }
diff --git a/test/expression/switch.stmt b/test/expression/switch.stmt
index 5d2d12cb..6d4aac4b 100644
--- a/test/expression/switch.stmt
+++ b/test/expression/switch.stmt
@@ -67,19 +67,11 @@ e = switch ([veryLongElement,veryLongElement,veryLongElement,]) {
   0 => "ok"
 };
 <<<
-### TODO(tall): Once block arguments are supported, this should become:
-### e = switch ([
-###   veryLongElement,
-###   veryLongElement,
-###   veryLongElement,
-### ]) {
-e = switch (
-  [
-    veryLongElement,
-    veryLongElement,
-    veryLongElement,
-  ]
-) {
+e = switch ([
+  veryLongElement,
+  veryLongElement,
+  veryLongElement,
+]) {
   0 => "ok",
 };
 >>> Split in case expression.
diff --git a/test/invocation/block_argument_comment.stmt b/test/invocation/block_argument_comment.stmt
new file mode 100644
index 00000000..a68cae64
--- /dev/null
+++ b/test/invocation/block_argument_comment.stmt
@@ -0,0 +1,107 @@
+40 columns                              |
+>>> Block comment before comma before block argument.
+function(first /* comment */, [second, third]);
+<<<
+function(first /* comment */, [
+  second,
+  third,
+]);
+>>> Block comment after comma before block argument.
+function(first, /* comment */ [second, third]);
+<<<
+function(first, /* comment */ [
+  second,
+  third,
+]);
+>>> Block comment inside block argument.
+function(first, [/* comment */ second, third]);
+<<<
+function(first, [
+  /* comment */ second,
+  third,
+]);
+>>>
+function(first, [second, third /* comment */]);
+<<<
+function(first, [
+  second,
+  third /* comment */,
+]);
+>>> Block comment after block argument.
+function([first, second] /* comment */, third);
+<<<
+function([
+  first,
+  second,
+] /* comment */, third);
+>>> Block comment after comma after block argument.
+function([first, second], /* comment */ third);
+<<<
+function([
+  first,
+  second,
+], /* comment */ third);
+>>> Block comment after argument after block argument.
+function([first, second], third /* comment */);
+<<<
+function([
+  first,
+  second,
+], third /* comment */);
+>>> Line comment before comma before block argument.
+function(first // comment
+, [second, third]);
+<<<
+function(
+  first, // comment
+  [second, third],
+);
+>>> Line comment after comma before block argument.
+function(first, // comment
+[second, third]);
+<<<
+function(
+  first, // comment
+  [second, third],
+);
+>>> Line comment inside block argument.
+function(first, [// comment
+second, third]);
+<<<
+function(first, [
+  // comment
+  second,
+  third,
+]);
+>>>
+function(first, [second, third // comment
+]);
+<<<
+function(first, [
+  second,
+  third, // comment
+]);
+>>> Line comment after block argument.
+function([first, second] // comment
+, third);
+<<<
+function(
+  [first, second], // comment
+  third,
+);
+>>> Line comment after comma after block argument.
+function([first, second], // comment
+third);
+<<<
+function(
+  [first, second], // comment
+  third,
+);
+>>> Line comment after argument after block argument.
+function([first, second], third // comment
+);
+<<<
+function(
+  [first, second],
+  third, // comment
+);
\ No newline at end of file
diff --git a/test/invocation/block_argument_kind.stmt b/test/invocation/block_argument_kind.stmt
new file mode 100644
index 00000000..d56c2118
--- /dev/null
+++ b/test/invocation/block_argument_kind.stmt
@@ -0,0 +1,284 @@
+40 columns                              |
+### Test which kinds of expressions are candidates to be block-formatted.
+>>> Non-empty block-bodied function expression.
+function(() { body; });
+<<<
+function(() {
+  body;
+});
+>>> Empty block-bodied function expression with a block comment.
+function(() { /* fairly long comment */ });
+<<<
+function(() {
+  /* fairly long comment */
+});
+>>> Empty block-bodied function expression with a line comment.
+function(() { // comment
+});
+<<<
+function(() {
+  // comment
+});
+>>> Block-bodied function expression with parameters.
+function((parameter) { body; });
+<<<
+function((parameter) {
+  body;
+});
+>>> Empty block-bodied function expression with parameters.
+function((parameter, anotherParameter) {});
+<<<
+function((
+  parameter,
+  anotherParameter,
+) {});
+>>> Expression-bodied function expression with parameters.
+function((parameter, anotherParameter) => body);
+<<<
+function((
+  parameter,
+  anotherParameter,
+) => body);
+>>> Block-bodied function expression with block comment in the parameters.
+function((/* very long block comment */) {});
+<<<
+function((
+  /* very long block comment */
+) {});
+>>> Block-bodied function expression with line comment in the parameters.
+function((// comment
+) {});
+<<<
+function((
+  // comment
+) {});
+>>> Expression-bodied function expression with block comment in the parameters.
+function((/* very long block comment */) => body);
+<<<
+function((
+  /* very long block comment */
+) => body);
+>>> Expression-bodied function expression with line comment in the parameters.
+function((// comment
+) => body);
+<<<
+function((
+  // comment
+) => body);
+>>> Empty block-bodied function expression with a block comment.
+function(() { /* fairly long comment */ });
+<<<
+function(() {
+  /* fairly long comment */
+});
+>>> Empty block-bodied function expression with a line comment.
+function(() { // comment
+});
+<<<
+function(() {
+  // comment
+});
+>>> An empty block-bodied function expression is not a block argument.
+function_________________________(() {});
+<<<
+function_________________________(
+  () {},
+);
+>>> A zero-parameter expression-bodied function is not a block argument.
+function_______________________(() => null);
+<<<
+function_______________________(
+  () => null,
+);
+>>> Function call.
+function(innerFunction(veryLongArgumentExpression));
+<<<
+function(innerFunction(
+  veryLongArgumentExpression,
+));
+>>> Zero-argument function call with block comment.
+function(innerFunction(/* long comment */));
+<<<
+function(innerFunction(
+  /* long comment */
+));
+>>> Zero-argument function call with line comment.
+function(innerFunction(// comment
+));
+<<<
+function(innerFunction(
+  // comment
+));
+>>> A zero-argument function call is not a block argument.
+function_______________________(inner());
+<<<
+function_______________________(
+  inner(),
+);
+>>> Method call.
+function(target.inner(veryLongArgumentExpression));
+<<<
+function(target.inner(
+  veryLongArgumentExpression,
+));
+>>> Zero-argument method call with block comment.
+function(target.inner(/* long comment */));
+<<<
+function(target.inner(
+  /* long comment */
+));
+>>> Zero-argument method call with line comment.
+function(target.inner(// comment
+));
+<<<
+function(target.inner(
+  // comment
+));
+>>> A zero-argument method call is not a block argument.
+function________________(target.inner());
+<<<
+function________________(
+  target.inner(),
+);
+>>> Instance creation expression.
+function(new SomeClass(veryLongArgumentExpression));
+<<<
+function(new SomeClass(
+  veryLongArgumentExpression,
+));
+>>> Zero-argument instance creation expression with block comment.
+function(new SomeClass(/* long comment */));
+<<<
+function(new SomeClass(
+  /* long comment */
+));
+>>> Zero-argument instance creation expression with line comment.
+function(new SomeClass(// comment
+));
+<<<
+function(new SomeClass(
+  // comment
+));
+>>> A zero-argument instance creation expression is not a block argument.
+function________________(new SomeClass());
+<<<
+function________________(
+  new SomeClass(),
+);
+>>> Parenthesized expression where inner expression is a block argument.
+function((innerFunction(veryLongArgumentExpression)));
+<<<
+function((innerFunction(
+  veryLongArgumentExpression,
+)));
+>>> List literal.
+function([veryLongElement, anotherLongElement]);
+<<<
+function([
+  veryLongElement,
+  anotherLongElement,
+]);
+>>> Empty list literal with block comment.
+function([/* a very long block comment */]);
+<<<
+function([
+  /* a very long block comment */
+]);
+>>> Empty list literal with line comment.
+function([// comment
+]);
+<<<
+function([
+  // comment
+]);
+>>> An empty list is not a block argument.
+function_____________________________([]);
+<<<
+function_____________________________(
+  [],
+);
+>>> Map literal.
+function({1: veryLongElement, 2: anotherLongElement});
+<<<
+function({
+  1: veryLongElement,
+  2: anotherLongElement,
+});
+>>> Set literal.
+function({veryLongElement, anotherLongElement});
+<<<
+function({
+  veryLongElement,
+  anotherLongElement,
+});
+>>> Empty set/map literal with block comment.
+function({/* a very long block comment */});
+<<<
+function({
+  /* a very long block comment */
+});
+>>> Empty set/map literal with line comment.
+function({// comment
+});
+<<<
+function({
+  // comment
+});
+>>> An empty map/set is not a block argument.
+function_____________________________({});
+<<<
+function_____________________________(
+  {},
+);
+>>> Record literal.
+function((veryLongElement, anotherLongElement));
+<<<
+function((
+  veryLongElement,
+  anotherLongElement,
+));
+>>> Empty record literal with block comment.
+function((/* a very long block comment */));
+<<<
+function((
+  /* a very long block comment */
+));
+>>> Empty record literal with line comment.
+function((// comment
+));
+<<<
+function((
+  // comment
+));
+>>> An empty record is not a block argument.
+function_____________________________(());
+<<<
+function_____________________________(
+  (),
+);
+>>> Switch expression.
+function(switch (n) {1 => veryLongElement, 2 => anotherLongElement});
+<<<
+function(switch (n) {
+  1 => veryLongElement,
+  2 => anotherLongElement,
+});
+>>> Empty switch expression with block comment.
+function(switch (n) {/* long comment */});
+<<<
+function(switch (n) {
+  /* long comment */
+});
+>>> Empty switch expression with line comment.
+function(switch (n) {// comment
+});
+<<<
+function(switch (n) {
+  // comment
+});
+>>> An empty switch expression is not a block argument.
+function___________________(switch (n) {});
+<<<
+function___________________(
+  switch (n) {},
+);
\ No newline at end of file
diff --git a/test/invocation/block_argument_multiple.stmt b/test/invocation/block_argument_multiple.stmt
new file mode 100644
index 00000000..894b1ab5
--- /dev/null
+++ b/test/invocation/block_argument_multiple.stmt
@@ -0,0 +1,66 @@
+40 columns                              |
+### Test how multiple block argument candidates are handled.
+>>> Multiple function expressions prevent block formatting.
+function(() { one; }, () { two; });
+<<<
+function(
+  () {
+    one;
+  },
+  () {
+    two;
+  },
+);
+>>> Empty and non-empty function expressions.
+function(() {}, () { body; }, () {});
+<<<
+function(() {}, () {
+  body;
+}, () {});
+>>> Function expression takes precedence over other kinds of block arguments.
+### The function is block formatted but the other arguments aren't.
+function([1, 2], () { body; }, {3, 4});
+<<<
+function([1, 2], () {
+  body;
+}, {3, 4});
+>>> Multiple collections prevent block formatting.
+function([element, element], {key: value});
+<<<
+function(
+  [element, element],
+  {key: value},
+);
+>>> Empty and non-empty collections.
+function([], [element, element], <String>{});
+<<<
+function([], [
+  element,
+  element,
+], <String>{});
+>>> Multiple calls prevent block formatting.
+function(inner(argument), new SomeClass(argument));
+<<<
+function(
+  inner(argument),
+  new SomeClass(argument),
+);
+>>> Empty and non-empty calls.
+function(a(), inner(argument), const C());
+<<<
+function(a(), inner(
+  argument,
+), const C());
+>>> Multiple switches prevent block formatting.
+function(switch (a) { 1 => 2 }, switch (b) { 1 => 2 });
+<<<
+function(
+  switch (a) { 1 => 2 },
+  switch (b) { 1 => 2 },
+);
+>>> Empty and non-empty switches.
+function(switch (a) {}, switch (b) { 1 => 2 }, switch (c) {});
+<<<
+function(switch (a) {}, switch (b) {
+  1 => 2,
+}, switch (c) {});
\ No newline at end of file
diff --git a/test/invocation/block_argument_other.stmt b/test/invocation/block_argument_other.stmt
new file mode 100644
index 00000000..a361081b
--- /dev/null
+++ b/test/invocation/block_argument_other.stmt
@@ -0,0 +1,89 @@
+40 columns                              |
+### Test other behavior related to block arguments.
+>>> Can have non-block arguments before the block argument.
+function(first, second, [third, fourth, fifth]);
+<<<
+function(first, second, [
+  third,
+  fourth,
+  fifth,
+]);
+>>> Can have non-block arguments after the block argument.
+function([first, second, third], fourth, fifth);
+<<<
+function([
+  first,
+  second,
+  third,
+], fourth, fifth);
+>>> Can have non-block arguments before and after the block argument.
+function(first, second, [third, fourth], fifth);
+<<<
+function(first, second, [
+  third,
+  fourth,
+], fifth);
+>>> Don't block argument format if it fits better to not.
+function('a long leading argument', <String>[element, element]);
+<<<
+function(
+  'a long leading argument',
+  <String>[element, element],
+);
+>>>
+function('a long leading argument',
+<String>[element, element, element, element]);
+<<<
+function(
+  'a long leading argument',
+  <String>[
+    element,
+    element,
+    element,
+    element,
+  ],
+);
+>>>
+function(<AVeryLong<Nested, GenericType>>[e]);
+<<<
+function(
+  <AVeryLong<Nested, GenericType>>[e],
+);
+>>>
+function(<AVeryLong<Nested, GenericType>>[element]);
+<<<
+function(
+  <AVeryLong<Nested, GenericType>>[
+    element,
+  ],
+);
+>>> A function block argument can be named.
+function(name: () { body; });
+<<<
+function(name: () {
+  body;
+});
+>>> A non-function block argument can be named.
+function(name: [element, element, element]);
+<<<
+function(name: [
+  element,
+  element,
+  element,
+]);
+>>> A long argument name can prevent block formatting.
+veryLongFunctionName(veryLongArgumentName: [element]);
+<<<
+veryLongFunctionName(
+  veryLongArgumentName: [element],
+);
+>>> A long argument name can prevent block formatting.
+veryLongFunctionName(veryLongArgumentName: [element, element, element]);
+<<<
+veryLongFunctionName(
+  veryLongArgumentName: [
+    element,
+    element,
+    element,
+  ],
+);
\ No newline at end of file
diff --git a/test/invocation/constructor.stmt b/test/invocation/constructor.stmt
index f57665b9..3c0d1264 100644
--- a/test/invocation/constructor.stmt
+++ b/test/invocation/constructor.stmt
@@ -47,4 +47,12 @@ new prefix.VeryLongClassName.veryLongNamedConstructor();
 <<<
 new prefix
     .VeryLongClassName
-    .veryLongNamedConstructor();
\ No newline at end of file
+    .veryLongNamedConstructor();
+>>> Allow block-formatted argument.
+new Future(new Duration(1), () {
+      print('I am a callback');
+    });
+<<<
+new Future(new Duration(1), () {
+  print('I am a callback');
+});
\ No newline at end of file
diff --git a/test/invocation/function.stmt b/test/invocation/function.stmt
index 869cb80c..a9a773c5 100644
--- a/test/invocation/function.stmt
+++ b/test/invocation/function.stmt
@@ -59,4 +59,25 @@ function<
     VeryLongTypeName,
     AnotherLongTypeName
   >
->(1, 2, 3);
\ No newline at end of file
+>(1, 2, 3);
+>>> Multiple nested split arguments.
+someFunctionOne(someArgument,
+someFunctionTwo(argument, argument, argument),
+someFunctionTwo(argument, argument, argument),
+someArgument, someArgument);
+<<<
+someFunctionOne(
+  someArgument,
+  someFunctionTwo(
+    argument,
+    argument,
+    argument,
+  ),
+  someFunctionTwo(
+    argument,
+    argument,
+    argument,
+  ),
+  someArgument,
+  someArgument,
+);
diff --git a/test/statement/assert.stmt b/test/statement/assert.stmt
index c7695eb1..d26fb942 100644
--- a/test/statement/assert.stmt
+++ b/test/statement/assert.stmt
@@ -49,10 +49,26 @@ assert(
 );
 <<<
 assert(1, 2);
->>> Add trailing comma if argument list split.
+>>> Add trailing comma if argument list splits.
 assert(longArgument1, veryLongArgument2);
 <<<
 assert(
   longArgument1,
   veryLongArgument2,
 );
+>>> Allow block formatting of condition.
+assert(inner(argument, argument, argument));
+<<<
+assert(inner(
+  argument,
+  argument,
+  argument,
+));
+>>> Allow block formatting of message.
+assert(true, inner(argument, argument, argument));
+<<<
+assert(true, inner(
+  argument,
+  argument,
+  argument,
+));
\ No newline at end of file
diff --git a/test/statement/switch.stmt b/test/statement/switch.stmt
index 11ba462b..8f99c84b 100644
--- a/test/statement/switch.stmt
+++ b/test/statement/switch.stmt
@@ -269,14 +269,10 @@ switch ([veryLongElement,veryLongElement]) {
     return "ok";
 }
 <<<
-### TODO(tall): Once block arguments are supported, this should become:
-### switch ([
-###   veryLongElement,
-###   veryLongElement,
-### ]) {
-switch (
-  [veryLongElement, veryLongElement]
-) {
+switch ([
+  veryLongElement,
+  veryLongElement,
+]) {
   case 0:
     return "ok";
 }
\ No newline at end of file
diff --git a/test/variable/local.stmt b/test/variable/local.stmt
index 38da4dcb..e7ebe61a 100644
--- a/test/variable/local.stmt
+++ b/test/variable/local.stmt
@@ -206,19 +206,11 @@ var variableName = switch (value) {
 var variableName = switch ([longElement, longElement, longElement])
 { 1 => 'one', 2 => 'two' };
 <<<
-### TODO(tall): Once block arguments are supported, this should become:
-### var variableName = switch ([
-###   longElement,
-###   longElement,
-###   longElement,
-### ]) {
-var variableName = switch (
-  [
-    longElement,
-    longElement,
-    longElement,
-  ]
-) {
+var variableName = switch ([
+  longElement,
+  longElement,
+  longElement,
+]) {
   1 => 'one',
   2 => 'two',
 };

From 7bf6dd4bdb9e59fb8c4b880cd807738c8cfcea90 Mon Sep 17 00:00:00 2001
From: Robert Nystrom <rnystrom@google.com>
Date: Thu, 16 Nov 2023 11:34:09 -0800
Subject: [PATCH 2/3] Apply review feedback.

Also renamed "isDelimited" (never a great name) to "canBlockSplit",
which I think better captures its use.
---
 lib/src/ast_extensions.dart                   | 53 ++++++++++---------
 lib/src/front_end/delimited_list_builder.dart |  4 +-
 lib/src/front_end/piece_factory.dart          |  2 +-
 lib/src/source_visitor.dart                   |  6 +--
 4 files changed, 34 insertions(+), 31 deletions(-)

diff --git a/lib/src/ast_extensions.dart b/lib/src/ast_extensions.dart
index 265e00cc..c63890a4 100644
--- a/lib/src/ast_extensions.dart
+++ b/lib/src/ast_extensions.dart
@@ -81,11 +81,11 @@ extension AstNodeExtensions on AstNode {
     if (node is SpreadElement) {
       var expression = node.expression;
       if (expression is ListLiteral) {
-        if (!expression.elements.isEmptyBody(expression.rightBracket)) {
+        if (expression.elements.canSplit(expression.rightBracket)) {
           return expression.leftBracket;
         }
       } else if (expression is SetOrMapLiteral) {
-        if (!expression.elements.isEmptyBody(expression.rightBracket)) {
+        if (expression.elements.canSplit(expression.rightBracket)) {
           return expression.leftBracket;
         }
       }
@@ -99,27 +99,27 @@ extension AstIterableExtensions on Iterable<AstNode> {
   /// Whether there is a comma token immediately following this.
   bool get hasCommaAfter => isNotEmpty && last.hasCommaAfter;
 
-  /// Whether the collection literal or block containing these nodes and
-  /// terminated by [rightBracket] is empty or not.
+  /// Whether the delimited construct containing these nodes and terminated by
+  /// [rightBracket] can have a split inside it.
   ///
-  /// An empty collection must have no elements or comments inside. Collections
-  /// like that are treated specially because they cannot be split inside.
-  bool isEmptyBody(Token rightBracket) =>
-      isEmpty && rightBracket.precedingComments == null;
+  /// We disallow splitting for entirely empty delimited constructs like `[]`,
+  /// but allow a split if there are elements or comments inside.
+  bool canSplit(Token rightBracket) =>
+      isNotEmpty || rightBracket.precedingComments != null;
 }
 
 extension ExpressionExtensions on Expression {
-  /// Whether this expression is a "delimited" one that allows block-like
-  /// formatting in some contexts. For example, in an assignment, a split in
-  /// the assigned value is usually indented:
+  /// Whether this expression is a non-empty delimited container for inner
+  /// expressions that allows "block-like" formatting in some contexts. For
+  /// example, in an assignment, a split in the assigned value is usually
+  /// indented:
   ///
   /// ```
   /// var variableName =
   ///     longValue;
   /// ```
   ///
-  /// But not if the initializer is a delimited expression and we don't split
-  /// at the `=`:
+  /// But if the initializer is block-like, we don't split at the `=`:
   ///
   /// ```
   /// var variableName = [
@@ -127,15 +127,18 @@ extension ExpressionExtensions on Expression {
   /// ];
   /// ```
   ///
-  /// Likewise, in an argument list, delimited expressions can be given
-  /// block-like formatting:
+  /// Likewise, in an argument list, block-like expressions can avoid splitting
+  /// the surrounding argument list:
   ///
   /// ```
   /// function([
   ///   element,
   /// ]);
   /// ```
-  bool get isDelimited {
+  ///
+  /// Completely empty delimited constructs like `[]` and `foo()` don't allow
+  /// splitting inside them, so are not considered block-like.
+  bool get canBlockSplit {
     // Unwrap named expressions to get the real expression inside.
     var expression = this;
     if (expression is NamedExpression) {
@@ -163,21 +166,21 @@ extension ExpressionExtensions on Expression {
     return switch (expression) {
       // A function expression can use either a non-empty parameter list or a
       // non-empty block body for block formatting.
-      FunctionExpression(:var parameters!, :var body) =>
-        !parameters.parameters.isEmptyBody(parameters.rightParenthesis) ||
-            body is BlockFunctionBody &&
-                !body.block.statements.isEmptyBody(body.block.rightBracket),
+      FunctionExpression(:var parameters?, :var body) =>
+        parameters.parameters.canSplit(parameters.rightParenthesis) ||
+            (body is BlockFunctionBody &&
+                body.block.statements.canSplit(body.block.rightBracket)),
       ListLiteral(:var elements, :var rightBracket) ||
       SetOrMapLiteral(:var elements, :var rightBracket) =>
-        !elements.isEmptyBody(rightBracket),
+        elements.canSplit(rightBracket),
       RecordLiteral(:var fields, :var rightParenthesis) =>
-        !fields.isEmptyBody(rightParenthesis),
+        fields.canSplit(rightParenthesis),
       SwitchExpression(:var cases, :var rightBracket) =>
-        !cases.isEmptyBody(rightBracket),
+        cases.canSplit(rightBracket),
       InstanceCreationExpression(:var argumentList) ||
       MethodInvocation(:var argumentList) =>
-        !argumentList.arguments.isEmptyBody(argumentList.rightParenthesis),
-      ParenthesizedExpression(:var expression) => expression.isDelimited,
+        argumentList.arguments.canSplit(argumentList.rightParenthesis),
+      ParenthesizedExpression(:var expression) => expression.canBlockSplit,
       _ => false,
     };
   }
diff --git a/lib/src/front_end/delimited_list_builder.dart b/lib/src/front_end/delimited_list_builder.dart
index b7940590..149a65ab 100644
--- a/lib/src/front_end/delimited_list_builder.dart
+++ b/lib/src/front_end/delimited_list_builder.dart
@@ -49,7 +49,7 @@ class DelimitedListBuilder {
   /// literal, etc.
   DelimitedListBuilder(this._visitor, [this._style = const ListStyle()]);
 
-  /// Create the final [ListPiece] out of the added brackets, delimiters,
+  /// Creates the final [ListPiece] out of the added brackets, delimiters,
   /// elements, and style.
   ListPiece build() {
     var blockElement = -1;
@@ -398,7 +398,7 @@ class DelimitedListBuilder {
 
     for (var i = 0; i < _blockCandidates.length; i++) {
       var node = _blockCandidates[i];
-      if (node is Expression && node.isDelimited) {
+      if (node is Expression && node.canBlockSplit) {
         if (node is FunctionExpression) {
           functions.add(i);
         } else {
diff --git a/lib/src/front_end/piece_factory.dart b/lib/src/front_end/piece_factory.dart
index a4e8f323..6a01bab8 100644
--- a/lib/src/front_end/piece_factory.dart
+++ b/lib/src/front_end/piece_factory.dart
@@ -611,7 +611,7 @@ mixin PieceFactory implements CommentWriter {
 
     var initializer = pieces.take();
     pieces.give(AssignPiece(target, initializer,
-        isValueDelimited: rightHandSide.isDelimited));
+        isValueDelimited: rightHandSide.canBlockSplit));
   }
 
   /// Finishes writing a named function declaration or anonymous function
diff --git a/lib/src/source_visitor.dart b/lib/src/source_visitor.dart
index 383a8fa7..221b57bf 100644
--- a/lib/src/source_visitor.dart
+++ b/lib/src/source_visitor.dart
@@ -364,7 +364,7 @@ class SourceVisitor extends ThrowingAstVisitor {
     // Treat empty blocks specially. In most cases, they are not allowed to
     // split. However, an empty block as the then statement of an if with an
     // else is always split.
-    if (node.statements.isEmptyBody(node.rightBracket)) {
+    if (!node.statements.canSplit(node.rightBracket)) {
       token(node.leftBracket);
       if (_splitEmptyBlock(node)) newline();
       token(node.rightBracket);
@@ -2803,7 +2803,7 @@ class SourceVisitor extends ThrowingAstVisitor {
 
   @override
   void visitSwitchExpression(SwitchExpression node) {
-    if (node.cases.isEmptyBody(node.rightBracket)) {
+    if (!node.cases.canSplit(node.rightBracket)) {
       // Don't allow splitting an empty switch expression.
       _visitSwitchValue(node.switchKeyword, node.leftParenthesis,
           node.expression, node.rightParenthesis);
@@ -4147,7 +4147,7 @@ class SourceVisitor extends ThrowingAstVisitor {
   /// Writes the brace-delimited body containing [nodes].
   void _visitBody(Token leftBracket, List<AstNode> nodes, Token rightBracket) {
     // Don't allow splitting in an empty body.
-    if (nodes.isEmptyBody(rightBracket)) {
+    if (!nodes.canSplit(rightBracket)) {
       token(leftBracket);
       token(rightBracket);
       return;

From 508409b6f4cbcd746adeafbfbd13c9cbbcffe64c Mon Sep 17 00:00:00 2001
From: Robert Nystrom <rnystrom@google.com>
Date: Thu, 16 Nov 2023 14:33:41 -0800
Subject: [PATCH 3/3] Remove _blockCandidates and store block element type on
 ListElement.

---
 lib/src/front_end/delimited_list_builder.dart | 33 +++++++++----------
 lib/src/piece/list.dart                       | 30 ++++++++++++++---
 2 files changed, 41 insertions(+), 22 deletions(-)

diff --git a/lib/src/front_end/delimited_list_builder.dart b/lib/src/front_end/delimited_list_builder.dart
index 149a65ab..ca3a6fc7 100644
--- a/lib/src/front_end/delimited_list_builder.dart
+++ b/lib/src/front_end/delimited_list_builder.dart
@@ -36,11 +36,6 @@ class DelimitedListBuilder {
 
   final ListStyle _style;
 
-  /// The nodes that have been visited by this builder.
-  ///
-  /// Used to find the block element from them, if desired.
-  final List<AstNode> _blockCandidates = [];
-
   /// The list of comments following the most recently written element before
   /// any comma following the element.
   CommentSequence _commentsBeforeComma = CommentSequence.empty;
@@ -131,8 +126,8 @@ class DelimitedListBuilder {
   /// [addCommentsBefore()] for the first token in the [piece].
   ///
   /// Assumes there is no comma after this piece.
-  void add(Piece piece) {
-    _elements.add(ListElement(piece));
+  void add(Piece piece, [BlockFormat format = BlockFormat.none]) {
+    _elements.add(ListElement(piece, format));
     _commentsBeforeComma = CommentSequence.empty;
   }
 
@@ -148,17 +143,21 @@ class DelimitedListBuilder {
     // Handle comments between the preceding element and this one.
     addCommentsBefore(element.firstNonCommentToken);
 
+    // See if it's an expression that supports block formatting.
+    var format = switch (element) {
+      FunctionExpression() when element.canBlockSplit => BlockFormat.function,
+      Expression() when element.canBlockSplit => BlockFormat.block,
+      _ => BlockFormat.none,
+    };
+
     // Traverse the element itself.
     _visitor.visit(element);
-    add(_visitor.pieces.split());
+    add(_visitor.pieces.split(), format);
 
     var nextToken = element.endToken.next!;
     if (nextToken.lexeme == ',') {
       _commentsBeforeComma = _visitor.takeCommentsBefore(nextToken);
     }
-
-    // If this element might be block formatted, remember it.
-    if (_style.allowBlockElement) _blockCandidates.add(element);
   }
 
   /// Inserts an inner left delimiter between two elements.
@@ -396,14 +395,14 @@ class DelimitedListBuilder {
     var functions = <int>[];
     var others = <int>[];
 
-    for (var i = 0; i < _blockCandidates.length; i++) {
-      var node = _blockCandidates[i];
-      if (node is Expression && node.canBlockSplit) {
-        if (node is FunctionExpression) {
+    for (var i = 0; i < _elements.length; i++) {
+      switch (_elements[i].blockFormat) {
+        case BlockFormat.function:
           functions.add(i);
-        } else {
+        case BlockFormat.block:
           others.add(i);
-        }
+        case BlockFormat.none:
+          break; // Not a block element.
       }
     }
 
diff --git a/lib/src/piece/list.dart b/lib/src/piece/list.dart
index 5713eb15..0455318c 100644
--- a/lib/src/piece/list.dart
+++ b/lib/src/piece/list.dart
@@ -187,6 +187,9 @@ class ListPiece extends Piece {
 final class ListElement {
   final Piece? _element;
 
+  /// What kind of block formatting can be applied to this element.
+  final BlockFormat blockFormat;
+
   /// If this piece has an opening delimiter after the comma, this is its
   /// lexeme, otherwise an empty string.
   ///
@@ -203,11 +206,14 @@ final class ListElement {
 
   final Piece? _comment;
 
-  ListElement(Piece element, [Piece? comment]) : this._(element, '', comment);
+  ListElement(Piece element, BlockFormat format, [Piece? comment])
+      : this._(element, format, '', comment);
 
-  ListElement.comment(Piece comment) : this._(null, '', comment);
+  ListElement.comment(Piece comment)
+      : this._(null, BlockFormat.none, '', comment);
 
-  ListElement._(this._element, this._delimiter, [this._comment]);
+  ListElement._(this._element, this.blockFormat, this._delimiter,
+      [this._comment]);
 
   /// Writes this element to [writer].
   ///
@@ -239,11 +245,11 @@ final class ListElement {
   /// Returns a new [ListElement] containing this one's element and [comment].
   ListElement withComment(Piece comment) {
     assert(_comment == null); // Shouldn't already have one.
-    return ListElement._(_element, _delimiter, comment);
+    return ListElement._(_element, blockFormat, _delimiter, comment);
   }
 
   ListElement withDelimiter(String delimiter) {
-    return ListElement._(_element, delimiter, _comment);
+    return ListElement._(_element, blockFormat, delimiter, _comment);
   }
 }
 
@@ -264,6 +270,20 @@ enum Commas {
   none,
 }
 
+/// What kind of block formatting style can be applied to the element.
+enum BlockFormat {
+  /// The element is a function expression, which takes priority over other
+  /// kinds of block formatted elements.
+  function,
+
+  /// The element is a collection literal or some other kind expression that
+  /// can be block formatted.
+  block,
+
+  /// The element can't be block formatted.
+  none,
+}
+
 /// The various ways a "list" can appear syntactically and be formatted.
 ///
 /// [ListPiece] is used for most places in code where a series of elements can