Refactor how child chunks are stored. (#1117)

* Refactor how child chunks are stored.

Given a piece of code like:

    main() {
      a();
      b();
    }

There are top level chunks for `main() }` and `}`, and the chunks for
`a()` and `b()` are children.

Previously, the code stored those child chunks in the preceding parent
chunk (so `main() {` here). But it's the subsequent chunk (`}`) that
determines whether the block contents are actually split, so that
doesn't make a lot of sense and leads to weird `+ 1` and `- 1` when
working with nested chunks.

This refactors the code so that child chunks are stored on the same
chunk that determines whether or not they split. This means that chunks
are now written in a post-order traversal: a block chunk's children come
before its own text.

Since we now know that a chunk will have children at the moment that
it's created, removed the old ChunkBlock class and replaced it with a
BlockChunk subclass of Chunk.

Also added a bunch of tests around trailing whitespace in blocks. When I
was testing this change on a corpus of code, I thought it inadvertently
changed some behavior, but it turns out that it was the previous
refactoring (which did deliberately change block formatting) and not
this one. These tests help pin that behavior down.

This commit here has zero changes on the formatter's visible behavior.

* Apply review feedback.
diff --git a/lib/src/chunk.dart b/lib/src/chunk.dart
index ec83adc..43e3d84 100644
--- a/lib/src/chunk.dart
+++ b/lib/src/chunk.dart
@@ -89,16 +89,6 @@
   ///             argument));
   final NestingLevel nesting;
 
-  /// If this chunk marks the beginning of a block, this contains the child
-  /// chunks and other data about that nested block.
-  ///
-  /// This should only be accessed when [isBlock] is `true`.
-  ChunkBlock get block => _block!;
-  ChunkBlock? _block;
-
-  /// Whether this chunk has a [block].
-  bool get isBlock => _block != null;
-
   /// The [Rule] that controls when a split should occur before this chunk.
   ///
   /// Multiple splits may share a [Rule].
@@ -136,16 +126,7 @@
   ///
   /// Does not include this chunk's own length, just the length of its child
   /// block chunks (recursively).
-  int get unsplitBlockLength {
-    if (!isBlock) return 0;
-
-    var length = 0;
-    for (var chunk in block.chunks) {
-      length += chunk.length + chunk.unsplitBlockLength;
-    }
-
-    return length;
-  }
+  int get unsplitBlockLength => 0;
 
   /// The [Span]s that contain this chunk.
   final spans = <Span>[];
@@ -189,29 +170,12 @@
     if (space != null) _spaceWhenUnsplit = space;
   }
 
-  /// Turns this chunk into one that can contain a block of child chunks.
-  void makeBlock(Chunk? blockArgument) {
-    assert(_block == null);
-    _block = ChunkBlock(blockArgument);
-  }
-
-  /// Returns `true` if the block body owned by this chunk should be expression
-  /// indented given a set of rule values provided by [getValue].
-  bool indentBlock(int Function(Rule) getValue) {
-    if (!isBlock) return false;
-
-    var argument = block.argument;
-    if (argument == null) return false;
-
-    var rule = argument.rule;
-
-    // There may be no rule if the block occurs inside a string interpolation.
-    // In that case, it's not clear if anything will look particularly nice, but
-    // expression nesting is probably marginally better.
-    if (rule == Rule.dummy) return true;
-
-    return rule.isSplit(getValue(rule), argument);
-  }
+  /// Returns `true` if this chunk is a block whose children should be
+  /// expression indented given a set of rule values provided by [getValue].
+  ///
+  /// [getValue] takes a [Rule] and returns the chosen split state value for
+  /// that [Rule].
+  bool indentBlock(int Function(Rule) getValue) => false;
 
   // Mark whether this chunk can divide the range of chunks.
   void markDivide(bool canDivide) {
@@ -234,9 +198,33 @@
   }
 }
 
-/// The child chunks owned by a chunk that begins a "block" -- an actual block
-/// statement, function expression, or collection literal.
-class ChunkBlock {
+/// A [Chunk] containing a list of nested "child" chunks that are formatted
+/// independently of the surrounding chunks.
+///
+/// This is used for blocks, function expressions, collection literals, etc.
+/// Basically, anywhere we have a delimited body of code whose formatting
+/// doesn't depend on how the surrounding code is formatted except to determine
+/// indentation.
+///
+/// This chunk's own text is the closing delimiter of the block, so its
+/// children come before itself. For example, given this code:
+///
+///     main() {
+///       var list = [
+///         element,
+///       ];
+///     }
+///
+/// It is organized into a tree of chunks like so:
+///
+///    - Chunk           "main() {"
+///    - BlockChunk
+///      |- Chunk          "var list = ["
+///      |- BlockChunk
+///      |  |- Chunk         "element,"
+///      |  '- (text)      "];"
+///      '- (text)       "}"
+class BlockChunk extends Chunk {
   /// If this block is for a collection literal in an argument list, this will
   /// be the chunk preceding this literal argument.
   ///
@@ -245,9 +233,39 @@
   final Chunk? argument;
 
   /// The child chunks in this block.
-  final List<Chunk> chunks = [];
+  final List<Chunk> children = [];
 
-  ChunkBlock(this.argument);
+  BlockChunk(this.argument, super.rule, super.indent, super.nesting,
+      {required super.space, required super.flushLeft})
+      : super(isDouble: false);
+
+  /// The unsplit length of all of this chunk's block contents.
+  ///
+  /// Does not include this chunk's own length, just the length of its child
+  /// block chunks (recursively).
+  @override
+  int get unsplitBlockLength {
+    var length = 0;
+    for (var chunk in children) {
+      length += chunk.length + chunk.unsplitBlockLength;
+    }
+
+    return length;
+  }
+
+  @override
+  bool indentBlock(int Function(Rule) getValue) {
+    var argument = this.argument;
+    if (argument == null) return false;
+
+    // There may be no rule if the block occurs inside a string interpolation.
+    // In that case, it's not clear if anything will look particularly nice, but
+    // expression nesting is probably marginally better.
+    var rule = argument.rule;
+    if (rule == Rule.dummy) return true;
+
+    return rule.isSplit(getValue(rule), argument);
+  }
 }
 
 /// The in-progress state for a [Span] that has been started but has not yet
diff --git a/lib/src/chunk_builder.dart b/lib/src/chunk_builder.dart
index a1bd230..e1eec94 100644
--- a/lib/src/chunk_builder.dart
+++ b/lib/src/chunk_builder.dart
@@ -567,28 +567,43 @@
     _blockArgumentNesting.removeLast();
   }
 
-  /// Starts a new block as a child of the current chunk.
+  /// Starts a new block chunk and returns the [ChunkBuilder] for it.
   ///
   /// Nested blocks are handled using their own independent [LineWriter].
-  ChunkBuilder startBlock([Chunk? argumentChunk]) {
-    var chunk = _chunks.last;
-    chunk.makeBlock(argumentChunk);
+  ChunkBuilder startBlock(
+      {Chunk? argumentChunk, bool indent = true, bool space = false}) {
+    // Start a block chunk for the block. It will contain the chunks for the
+    // contents of the block, and its own text will be the closing block
+    // delimiter.
+    var chunk = BlockChunk(argumentChunk, _rules.last, _nesting.indentation,
+        _blockArgumentNesting.last,
+        space: space, flushLeft: _pendingFlushLeft);
+    _chunks.add(chunk);
+    _pendingFlushLeft = false;
 
-    return ChunkBuilder._(this, _formatter, _source, chunk.block.chunks);
+    var builder = ChunkBuilder._(this, _formatter, _source, chunk.children);
+
+    if (indent) builder.indent();
+
+    // Create a hard split for the contents. The rule on the parent BlockChunk
+    // determines whether the body is split or not. This hard rule is only when
+    // the block's contents are split.
+    var rule = Rule.hard();
+    builder.startRule(rule);
+    builder.split(nest: false, space: space);
+
+    return builder;
   }
 
   /// Ends this [ChunkBuilder], which must have been created by [startBlock()].
   ///
   /// Forces the chunk that owns the block to split if it can tell that the
   /// block contents will always split. It does that by looking for hard splits
-  /// in the block. If [bodyRule] is given, that rule will be ignored when
-  /// determining if a block contains a hard split. If [space] is `true`, the
-  /// split at the end of the block will get a space when unsplit. If
+  /// in the block that aren't for top level elements in the block. If
   /// [forceSplit] is `true`, the block always splits.
   ///
   /// Returns the previous writer for the surrounding block.
-  ChunkBuilder endBlock(
-      {Rule? bodyRule, bool space = false, bool forceSplit = true}) {
+  ChunkBuilder endBlock({bool forceSplit = true}) {
     _divideChunks();
 
     // If the last chunk ends with a comment that wants a newline after it,
@@ -606,37 +621,21 @@
           break;
         }
 
-        // If there are any hardened splits in the chunks (aside from the first
-        // one which is always a hard split since it is the beginning of the
-        // code), then force the collection to split.
-        if (chunk != _chunks.first &&
-            chunk.rule.isHardened &&
-            chunk.rule != bodyRule) {
+        // If there are any hardened splits in the chunks (aside from ones
+        // using the initial hard rule created by [startBlock()] which are for
+        // the top level elements in the block), then force the block to split.
+        if (chunk.rule.isHardened && chunk.rule != _rules.first) {
           forceSplit = true;
           break;
         }
       }
     }
 
-    if (forceSplit) forceRules();
-
-    var parent = _parent!;
-    parent._endChildBlock(space: space, forceSplit: forceSplit);
-    return parent;
-  }
-
-  /// Finishes off the last chunk in a child block of this parent.
-  void _endChildBlock({required bool space, required bool forceSplit}) {
     // If there is a hard newline within the block, force the surrounding rule
     // for it so that we apply that constraint.
-    if (forceSplit) forceRules();
-
-    // Start a new chunk for the code after the block contents. The split at
-    // the beginning of this chunk also determines whether the preceding block
-    // splits and, if so, how it is indented.
-    _startChunk(_blockArgumentNesting.last, isHard: false, space: space);
-
-    if (rule.isHardened) _handleHardSplit();
+    var parent = _parent!;
+    if (forceSplit) parent.forceRules();
+    return parent;
   }
 
   /// Finishes writing and returns a [SourceCode] containing the final output
@@ -902,9 +901,9 @@
     if (!chunk.rule.isHardened) return false;
     if (chunk.nesting.isNested) return false;
 
-    // If the previous chunk is a block, then this chunk determines whether the
-    // block contents split, so don't separate it from the block.
-    if (i > 0 && _chunks[i - 1].isBlock) return false;
+    // If the chunk is the ending delimiter of a block, then don't separate it
+    // and its children from the preceding beginning of the block.
+    if (_chunks[i] is BlockChunk) return false;
 
     return true;
   }
diff --git a/lib/src/debug.dart b/lib/src/debug.dart
index 567e442..351e80e 100644
--- a/lib/src/debug.dart
+++ b/lib/src/debug.dart
@@ -70,7 +70,7 @@
     for (var chunk in chunks) {
       spanSet.addAll(chunk.spans);
 
-      if (chunk.isBlock) addSpans(chunk.block.chunks);
+      if (chunk is BlockChunk) addSpans(chunk.children);
     }
   }
 
@@ -82,11 +82,17 @@
   var rows = <List<String>>[];
 
   void addChunk(List<Chunk> chunks, String prefix, int index) {
+    var chunk = chunks[index];
+
+    if (chunk is BlockChunk) {
+      for (var j = 0; j < chunk.children.length; j++) {
+        addChunk(chunk.children, '$prefix$index.', j);
+      }
+    }
+
     var row = <String>[];
     row.add('$prefix$index:');
 
-    var chunk = chunks[index];
-
     void writeIf(predicate, String Function() callback) {
       if (predicate) {
         row.add(callback());
@@ -148,12 +154,6 @@
     }
 
     rows.add(row);
-
-    if (chunk.isBlock) {
-      for (var j = 0; j < chunk.block.chunks.length; j++) {
-        addChunk(chunk.block.chunks, '$prefix$index.', j);
-      }
-    }
   }
 
   for (var i = start; i < chunks.length; i++) {
@@ -222,10 +222,11 @@
   void writeChunksUnsplit(List<Chunk> chunks) {
     for (var chunk in chunks) {
       if (chunk.spaceWhenUnsplit) buffer.write(' ');
-      buffer.write(chunk.text);
 
       // Recurse into the block.
-      if (chunk.isBlock) writeChunksUnsplit(chunk.block.chunks);
+      if (chunk is BlockChunk) writeChunksUnsplit(chunk.children);
+
+      buffer.write(chunk.text);
     }
   }
 
@@ -241,11 +242,11 @@
       buffer.write(' ');
     }
 
-    buffer.write(chunk.text);
-
-    if (chunk.isBlock && !splits.shouldSplitAt(i)) {
-      writeChunksUnsplit(chunk.block.chunks);
+    if (chunk is BlockChunk && !splits.shouldSplitAt(i)) {
+      writeChunksUnsplit(chunk.children);
     }
+
+    buffer.write(chunk.text);
   }
 
   log(buffer);
diff --git a/lib/src/line_splitting/solve_state.dart b/lib/src/line_splitting/solve_state.dart
index 8fb320b..90abec7 100644
--- a/lib/src/line_splitting/solve_state.dart
+++ b/lib/src/line_splitting/solve_state.dart
@@ -297,7 +297,7 @@
           // And any expression nesting.
           indent += chunk.nesting.totalUsedIndent;
 
-          if (i > 0 && _splitter.chunks[i - 1].indentBlock(getValue)) {
+          if (_splitter.chunks[i].indentBlock(getValue)) {
             indent += Indent.expression;
           }
         }
@@ -389,19 +389,18 @@
         if (chunk.spaceWhenUnsplit) length++;
       }
 
-      length += chunk.text.length;
-
-      if (chunk.isBlock) {
-        if (_splits.shouldSplitAt(i + 1)) {
+      if (chunk is BlockChunk) {
+        if (_splits.shouldSplitAt(i)) {
           // Include the cost of the nested block.
-          cost += _splitter.writer
-              .formatBlock(chunk, _splits.getColumn(i + 1))
-              .cost;
+          cost +=
+              _splitter.writer.formatBlock(chunk, _splits.getColumn(i)).cost;
         } else {
           // Include the nested block inline, if any.
           length += chunk.unsplitBlockLength;
         }
       }
+
+      length += chunk.text.length;
     }
 
     // Add the costs for the rules that have any splits.
diff --git a/lib/src/line_writer.dart b/lib/src/line_writer.dart
index 31ecdd2..ae5866c 100644
--- a/lib/src/line_writer.dart
+++ b/lib/src/line_writer.dart
@@ -67,7 +67,7 @@
   ///     }
   ///
   /// When we format the function expression's body, [column] will be 2, not 4.
-  FormatResult formatBlock(Chunk chunk, int column) {
+  FormatResult formatBlock(BlockChunk chunk, int column) {
     var key = _BlockKey(chunk, column);
 
     // Use the cached one if we have it.
@@ -75,7 +75,7 @@
     if (cached != null) return cached;
 
     var writer = LineWriter._(
-        chunk.block.chunks, _lineEnding, pageWidth, column, _blockCache);
+        chunk.children, _lineEnding, pageWidth, column, _blockCache);
     return _blockCache[key] = writer.writeLines();
   }
 
@@ -128,6 +128,32 @@
     for (var i = 0; i < chunks.length; i++) {
       var chunk = chunks[i];
 
+      // Write the block chunk's children first.
+      if (chunk is BlockChunk) {
+        if (!splits.shouldSplitAt(i)) {
+          // This block didn't split (which implies none of the child blocks
+          // of that block split either, recursively), so write them all inline.
+          _writeChunksUnsplit(chunk);
+        } else {
+          _buffer.write(_lineEnding);
+
+          // Include the formatted block contents.
+          var block = formatBlock(chunk, splits.getColumn(i));
+
+          // If this block contains one of the selection markers, tell the
+          // writer where it ended up in the final output.
+          if (block.selectionStart != null) {
+            _selectionStart = length + block.selectionStart!;
+          }
+
+          if (block.selectionEnd != null) {
+            _selectionEnd = length + block.selectionEnd!;
+          }
+
+          _buffer.write(block.text);
+        }
+      }
+
       if (splits.shouldSplitAt(i)) {
         // Don't write an initial single newline at the beginning of the output.
         // If this is for a block, then the newline will be written before
@@ -144,49 +170,20 @@
       }
 
       _writeChunk(chunk);
-
-      if (chunk.isBlock) {
-        // The block children of this chunk are followed by the next chunk. If
-        // the block splits, then we will also split before that next chunk.
-        // That means the chunk after this one determines whether the block
-        // contents are split.
-        if (!splits.shouldSplitAt(i + 1)) {
-          // This block didn't split (which implies none of the child blocks
-          // of that block split either, recursively), so write them all inline.
-          _writeChunksUnsplit(chunk);
-        } else {
-          _buffer.write(_lineEnding);
-
-          // Include the formatted block contents.
-          var block = formatBlock(chunk, splits.getColumn(i + 1));
-
-          // If this block contains one of the selection markers, tell the
-          // writer where it ended up in the final output.
-          if (block.selectionStart != null) {
-            _selectionStart = length + block.selectionStart!;
-          }
-
-          if (block.selectionEnd != null) {
-            _selectionEnd = length + block.selectionEnd!;
-          }
-
-          _buffer.write(block.text);
-        }
-      }
     }
 
     return splits.cost;
   }
 
-  /// Writes the block chunks of [chunk] (and any child chunks of them,
+  /// Writes the block chunks of [block] (and any child chunks of them,
   /// recursively) without any splitting.
-  void _writeChunksUnsplit(Chunk chunk) {
-    for (var blockChunk in chunk.block.chunks) {
-      if (blockChunk.spaceWhenUnsplit) _buffer.write(' ');
-      _writeChunk(blockChunk);
+  void _writeChunksUnsplit(BlockChunk block) {
+    for (var chunk in block.children) {
+      if (chunk.spaceWhenUnsplit) _buffer.write(' ');
+      _writeChunk(chunk);
 
       // Recurse into the block.
-      if (blockChunk.isBlock) _writeChunksUnsplit(blockChunk);
+      if (chunk is BlockChunk) _writeChunksUnsplit(chunk);
     }
   }
 
diff --git a/lib/src/source_visitor.dart b/lib/src/source_visitor.dart
index 2810ebc..8e5b6c4 100644
--- a/lib/src/source_visitor.dart
+++ b/lib/src/source_visitor.dart
@@ -541,7 +541,7 @@
     // Process the inner cascade sections as a separate block. This way the
     // entire cascade expression isn't line split as a single monolithic unit,
     // which is very slow.
-    builder = builder.startBlock();
+    builder = builder.startBlock(indent: false);
 
     for (var i = 0; i < node.cascadeSections.length - 1; i++) {
       newline();
@@ -1022,7 +1022,7 @@
 
     builder.unnest();
 
-    var rule = _beginBody(node.leftBracket, space: true);
+    _beginBody(node.leftBracket, space: true);
 
     visitCommaSeparatedNodes(node.constants, between: splitOrTwoNewlines);
 
@@ -1058,10 +1058,7 @@
 
     _visitBodyContents(node.members);
 
-    builder.endRule();
-
-    _endBody(rule, node.rightBracket,
-        space: true,
+    _endBody(node.rightBracket,
         forceSplit: semicolon != null ||
             trailingComma != null ||
             node.members.isNotEmpty);
@@ -2464,9 +2461,9 @@
     space();
     builder.unnest();
 
-    var rule = _beginBody(node.leftBracket);
+    _beginBody(node.leftBracket);
     visitNodes(node.members, between: oneOrTwoNewlines, after: newline);
-    _endBody(rule, node.rightBracket, forceSplit: true);
+    _endBody(node.rightBracket, forceSplit: true);
   }
 
   @override
@@ -2994,7 +2991,7 @@
     // Add this collection to the stack.
     _collectionSplits.add(false);
 
-    var rule = _beginBody(leftBracket);
+    _beginBody(leftBracket);
     if (node != null) _startPossibleConstContext(node.constKeyword);
 
     // If a collection contains a line comment, we assume it's a big complex
@@ -3035,8 +3032,6 @@
       }
     }
 
-    builder.endRule();
-
     // If there is a collection inside this one, it forces this one to split.
     var force = _collectionSplits.removeLast();
 
@@ -3044,7 +3039,7 @@
     if (elements.hasCommaAfter) force = true;
 
     if (node != null) _endPossibleConstContext(node.constKeyword);
-    _endBody(rule, rightBracket, forceSplit: force);
+    _endBody(rightBracket, forceSplit: force);
   }
 
   /// Writes [parameters], which is assumed to have a trailing comma after the
@@ -3084,7 +3079,6 @@
 
     // Process the parameters as a separate set of chunks.
     builder = builder.startBlock();
-    builder.indent();
 
     for (var parameter in parameters.parameters) {
       newline();
@@ -3279,15 +3273,10 @@
   /// If [space] is `true`, writes a space after [leftBracket] when not split.
   ///
   /// Writes the delimiter (with a space after it when unsplit if [space] is
-  /// `true`), then creates and returns the [Rule] that handles splitting the
-  /// body.
-  Rule _beginBody(Token leftBracket, {bool space = false}) {
+  /// `true`).
+  void _beginBody(Token leftBracket, {bool space = false}) {
     token(leftBracket);
 
-    // TODO(rnystrom): This rule is only used for the chunk with the closing
-    // delimiter since the literal body child chunks have their own rule. Is
-    // there a cleaner way to handle this?
-
     // Create a rule for whether or not to split the block contents. If this
     // literal is associated with an argument list or if element that wants to
     // handle splitting and indenting it, use its rule. Otherwise, use a
@@ -3295,32 +3284,19 @@
     builder.startRule(_blockRules[leftBracket]);
 
     // Process the contents as a separate set of chunks.
-    builder = builder.startBlock(_blockPreviousChunks[leftBracket]);
-    builder.indent();
-
-    // Create a hard split for the contents. The parent chunk of the body
-    // handles the unsplit case, so this only comes into play when the body
-    // splits.
-    var rule = Rule.hard();
-    builder.startRule(rule);
-    builder.split(nest: false, space: space);
-    return rule;
+    builder = builder.startBlock(
+        argumentChunk: _blockPreviousChunks[leftBracket], space: space);
   }
 
   /// Ends the body started by a call to [_beginBody()].
   ///
-  /// [bodyRule] is the [Rule] returned by the previous call to [_beginBody()].
   /// If [space] is `true`, writes a space before the closing bracket when not
   /// split. If [forceSplit] is `true`, forces the body to split.
-  void _endBody(Rule bodyRule, Token rightBracket,
-      {bool space = false, bool forceSplit = false}) {
+  void _endBody(Token rightBracket, {bool forceSplit = false}) {
     // Put comments before the closing delimiter inside the block.
     var hasLeadingNewline = writePrecedingCommentsAndNewlines(rightBracket);
 
-    builder = builder.endBlock(
-        bodyRule: bodyRule,
-        space: space,
-        forceSplit: hasLeadingNewline || forceSplit);
+    builder = builder.endBlock(forceSplit: hasLeadingNewline || forceSplit);
 
     builder.endRule();
 
@@ -3413,9 +3389,9 @@
       return;
     }
 
-    var rule = _beginBody(leftBracket);
+    _beginBody(leftBracket);
     _visitBodyContents(nodes);
-    _endBody(rule, rightBracket, forceSplit: nodes.isNotEmpty);
+    _endBody(rightBracket, forceSplit: nodes.isNotEmpty);
   }
 
   /// Writes the string literal [string] to the output.
diff --git a/test/comments/classes.unit b/test/comments/classes.unit
index 8f65e86..527e75a 100644
--- a/test/comments/classes.unit
+++ b/test/comments/classes.unit
@@ -107,7 +107,7 @@
   /// doc
   var b = 2;
 }
->>> remove blank line before beginning of body
+>>> remove blank lines before beginning of body
 class A {
 
 
@@ -118,6 +118,17 @@
 class A {
   // comment
 }
+>>> remove blank lines after end of body
+class A {
+  // comment
+
+
+
+}
+<<<
+class A {
+  // comment
+}
 >>> nested flush left comment
 class Foo {
   method() {
diff --git a/test/comments/enums.unit b/test/comments/enums.unit
index 28df90d..819414d 100644
--- a/test/comments/enums.unit
+++ b/test/comments/enums.unit
@@ -39,7 +39,7 @@
 
   /* comment */
 }
->>> remove blank line before beginning of body
+>>> remove blank lines before beginning of body
 enum A {
 
 
@@ -52,6 +52,19 @@
   // comment
   B
 }
+>>> remove blank lines after end of body
+enum A {
+  B
+  // comment
+
+
+
+}
+<<<
+enum A {
+  B
+  // comment
+}
 >>> ensure blank line above doc comments
 enum Foo {/// doc
 a,/// doc
diff --git a/test/comments/extensions.unit b/test/comments/extensions.unit
index 67123f3..007317f 100644
--- a/test/comments/extensions.unit
+++ b/test/comments/extensions.unit
@@ -126,7 +126,7 @@
   /// doc
   b() => 2;
 }
->>> remove blank line before beginning of body
+>>> remove blank lines before beginning of body
 extension A on B {
 
 
@@ -137,6 +137,17 @@
 extension A on B {
   // comment
 }
+>>> remove blank lines after end of body
+extension A on B {
+  // comment
+
+
+
+}
+<<<
+extension A on B {
+  // comment
+}
 >>> nested flush left comment
 extension A on B {
   method() {
diff --git a/test/comments/functions.unit b/test/comments/functions.unit
index a1288e9..0443087 100644
--- a/test/comments/functions.unit
+++ b/test/comments/functions.unit
@@ -116,7 +116,7 @@
 <<<
 longFunction(
     /* a very long block comment */) {}
->>> remove blank line before beginning of body
+>>> remove blank lines before beginning of body
 main() {
 
 
@@ -127,6 +127,17 @@
 main() {
   // comment
 }
+>>> remove blank lines after end of body
+main() {
+  // comment
+
+
+
+}
+<<<
+main() {
+  // comment
+}
 >>> comment before "]" with trailing comma
 function([parameter,/* c */]) {;}
 <<<
diff --git a/test/comments/lists.stmt b/test/comments/lists.stmt
index 0f7aa04..09269a4 100644
--- a/test/comments/lists.stmt
+++ b/test/comments/lists.stmt
@@ -89,7 +89,7 @@
 var list = [1,/* a */ 2 /* b */  , 3];
 <<<
 var list = [1, /* a */ 2 /* b */, 3];
->>> remove blank line before beginning of body
+>>> remove blank lines before beginning of body
 var list = [
 
 
@@ -99,4 +99,15 @@
 <<<
 var list = [
   // comment
+];
+>>> remove blank lines after end of body
+var list = [
+  // comment
+
+
+
+];
+<<<
+var list = [
+  // comment
 ];
\ No newline at end of file
diff --git a/test/comments/maps.stmt b/test/comments/maps.stmt
index a4125f5..cd7888d 100644
--- a/test/comments/maps.stmt
+++ b/test/comments/maps.stmt
@@ -94,7 +94,7 @@
   // comment
   'foo': 1
 };
->>> remove blank line before beginning of body
+>>> remove blank lines before beginning of body
 var map = {
 
 
@@ -104,4 +104,15 @@
 <<<
 var map = {
   // comment
+};
+>>> remove blank lines after end of body
+var map = {
+  // comment
+
+
+
+};
+<<<
+var map = {
+  // comment
 };
\ No newline at end of file
diff --git a/test/comments/sets.stmt b/test/comments/sets.stmt
index 5ec54a7..79aa505 100644
--- a/test/comments/sets.stmt
+++ b/test/comments/sets.stmt
@@ -106,4 +106,15 @@
 <<<
 var set = <int>{
   // comment
+};
+>>> remove blank lines after end of body
+var set = <int>{
+  // comment
+
+
+
+};
+<<<
+var set = <int>{
+  // comment
 };
\ No newline at end of file
diff --git a/test/comments/statements.stmt b/test/comments/statements.stmt
index f196291..2eadaad 100644
--- a/test/comments/statements.stmt
+++ b/test/comments/statements.stmt
@@ -44,7 +44,7 @@
 /*
 */
 var i = value;
->>> remove blank line before beginning of block
+>>> remove blank lines before beginning of block
 while (true) {
 
 
@@ -55,6 +55,17 @@
 while (true) {
   // comment
 }
+>>> remove blank lines after end of block
+while (true) {
+  // comment
+
+
+
+}
+<<<
+while (true) {
+  // comment
+}
 >>>
 main() {
   /* comment */ statement;
diff --git a/test/whitespace/classes.unit b/test/whitespace/classes.unit
index a4ea79f..15841bc 100644
--- a/test/whitespace/classes.unit
+++ b/test/whitespace/classes.unit
@@ -156,4 +156,17 @@
   abstract covariant var a, b;
   abstract final int c;
   abstract int i;
+}
+>>> discard trailing newlines in body
+class Foo {
+  bar() {}
+
+
+
+
+
+}
+<<<
+class Foo {
+  bar() {}
 }
\ No newline at end of file
diff --git a/test/whitespace/enums.unit b/test/whitespace/enums.unit
index dc76c6b..17107ce 100644
--- a/test/whitespace/enums.unit
+++ b/test/whitespace/enums.unit
@@ -280,3 +280,35 @@
   c(123),
   ;
 }
+>>> discard trailing newlines in body
+enum E {
+  a,
+
+
+
+
+
+}
+<<<
+enum E {
+  a,
+}
+>>> discard trailing newlines in body
+enum E {
+  a;
+
+
+
+  bar() {}
+
+
+
+
+
+}
+<<<
+enum E {
+  a;
+
+  bar() {}
+}
\ No newline at end of file
diff --git a/test/whitespace/extensions.unit b/test/whitespace/extensions.unit
index 009bbdf..1753009 100644
--- a/test/whitespace/extensions.unit
+++ b/test/whitespace/extensions.unit
@@ -117,4 +117,17 @@
 >>> unnamed generic extension
 extension   <  T  ,  S  >    on   B{}
 <<<
-extension<T, S> on B {}
\ No newline at end of file
+extension<T, S> on B {}
+>>> discard trailing newlines in body
+extension A on B {
+  bar() {}
+
+
+
+
+
+}
+<<<
+extension A on B {
+  bar() {}
+}
\ No newline at end of file
diff --git a/test/whitespace/methods.unit b/test/whitespace/methods.unit
index a27c443..9d7b687 100644
--- a/test/whitespace/methods.unit
+++ b/test/whitespace/methods.unit
@@ -56,4 +56,20 @@
 <<<
 class A {
   A(covariant this.foo);
+}
+>>> discard trailing newlines in method body
+class Foo {
+  bar() {
+    baz();
+
+
+
+
+  }
+}
+<<<
+class Foo {
+  bar() {
+    baz();
+  }
 }
\ No newline at end of file